modelsimpl.h 1014 Bytes
Newer Older
1
#pragma once
Shahriar's avatar
Shahriar committed
2
3
4

#include <torch/torch.h>

5
6
7
8
#ifndef TORCH_CHECK
#define TORCH_CHECK AT_CHECK
#endif

Shahriar's avatar
Shahriar committed
9
10
11
12
13
14
15
namespace vision {
namespace models {
namespace modelsimpl {

// TODO here torch::relu_ and torch::adaptive_avg_pool2d wrapped in
// torch::nn::Fuctional don't work. so keeping these for now

16
17
inline torch::Tensor& relu_(const torch::Tensor& x) {
  return x.relu_();
Shahriar's avatar
Shahriar committed
18
19
}

20
inline torch::Tensor& relu6_(const torch::Tensor& x) {
21
  return x.clamp_(0, 6);
Shahriar's avatar
Shahriar committed
22
23
24
}

inline torch::Tensor adaptive_avg_pool2d(
25
    const torch::Tensor& x,
Shahriar's avatar
Shahriar committed
26
27
28
29
30
    torch::ExpandingArray<2> output_size) {
  return torch::adaptive_avg_pool2d(x, output_size);
}

inline torch::Tensor max_pool2d(
31
    const torch::Tensor& x,
Shahriar's avatar
Shahriar committed
32
33
34
35
36
37
38
39
40
41
42
43
    torch::ExpandingArray<2> kernel_size,
    torch::ExpandingArray<2> stride) {
  return torch::max_pool2d(x, kernel_size, stride);
}

inline bool double_compare(double a, double b) {
  return double(std::abs(a - b)) < std::numeric_limits<double>::epsilon();
};

} // namespace modelsimpl
} // namespace models
} // namespace vision