kSpatialDim Class — pytorch Architecture

Source Code

aten/src/ATen/native/quantized/cudnn/Pooling.cpp lines 223–237

template <uint32_t kSpatialDim>
class QMaxPool_arr_args final {
 public:
  static Tensor run(
      const Tensor& qx,
      std::vector<int64_t> kernel_size,
      std::vector<int64_t> stride,
      std::vector<int64_t> padding,
      std::vector<int64_t> dilation,
      bool ceil_mode) {
    static_assert(kSpatialDim == 2, "quantized max pool is only valid for 2D");
    return quantized_max_pool2d_cudnn(qx, kernel_size, stride, padding,
                                    dilation, ceil_mode);
  }
};