Assertion macros compatible with pytorch master (#540)

Author: ahirner

torchvision/csrc/cpu/ROIAlign_cpu.cpp

@@ -223,8 +223,8 @@ at::Tensor ROIAlign_forward_cpu(const at::Tensor& input,
                                 const int pooled_height,
                                 const int pooled_width,
                                 const int sampling_ratio) {
-  AT_ASSERT(!input.type().is_cuda(), "input must be a CPU tensor");
-  AT_ASSERT(!rois.type().is_cuda(), "rois must be a CPU tensor");
+  AT_ASSERTM(!input.type().is_cuda(), "input must be a CPU tensor");
+  AT_ASSERTM(!rois.type().is_cuda(), "rois must be a CPU tensor");
 
   auto num_rois = rois.size(0);
   auto channels = input.size(1);

torchvision/csrc/cpu/nms_cpu.cpp

@@ -5,9 +5,9 @@ template <typename scalar_t>
 at::Tensor nms_cpu_kernel(const at::Tensor& dets,
                           const at::Tensor& scores,
                           const float threshold) {
-  AT_ASSERT(!dets.type().is_cuda(), "dets must be a CPU tensor");
-  AT_ASSERT(!scores.type().is_cuda(), "scores must be a CPU tensor");
-  AT_ASSERT(dets.type() == scores.type(), "dets should have the same type as scores");
+  AT_ASSERTM(!dets.type().is_cuda(), "dets must be a CPU tensor");
+  AT_ASSERTM(!scores.type().is_cuda(), "scores must be a CPU tensor");
+  AT_ASSERTM(dets.type() == scores.type(), "dets should have the same type as scores");
 
   if (dets.numel() == 0)
     return torch::CPU(at::kLong).tensor();

torchvision/csrc/cuda/ROIAlign_cuda.cu

@@ -258,8 +258,8 @@ at::Tensor ROIAlign_forward_cuda(const at::Tensor& input,
                                  const int pooled_height,
                                  const int pooled_width,
                                  const int sampling_ratio) {
-  AT_ASSERT(input.type().is_cuda(), "input must be a CUDA tensor");
-  AT_ASSERT(rois.type().is_cuda(), "rois must be a CUDA tensor");
+  AT_ASSERTM(input.type().is_cuda(), "input must be a CUDA tensor");
+  AT_ASSERTM(rois.type().is_cuda(), "rois must be a CUDA tensor");
 
   auto num_rois = rois.size(0);
   auto channels = input.size(1);
@@ -308,8 +308,8 @@ at::Tensor ROIAlign_backward_cuda(const at::Tensor& grad,
                                   const int height,
                                   const int width,
                                   const int sampling_ratio) {
-  AT_ASSERT(grad.type().is_cuda(), "grad must be a CUDA tensor");
-  AT_ASSERT(rois.type().is_cuda(), "rois must be a CUDA tensor");
+  AT_ASSERTM(grad.type().is_cuda(), "grad must be a CUDA tensor");
+  AT_ASSERTM(rois.type().is_cuda(), "rois must be a CUDA tensor");
 
   auto num_rois = rois.size(0);
   at::Tensor grad_input = grad.type().tensor({batch_size, channels, height, width}).zero_();

torchvision/csrc/cuda/ROIPool_cuda.cu

@@ -110,8 +110,8 @@ std::tuple<at::Tensor, at::Tensor> ROIPool_forward_cuda(const at::Tensor& input,
                                 const float spatial_scale,
                                 const int pooled_height,
                                 const int pooled_width) {
-  AT_ASSERT(input.type().is_cuda(), "input must be a CUDA tensor");
-  AT_ASSERT(rois.type().is_cuda(), "rois must be a CUDA tensor");
+  AT_ASSERTM(input.type().is_cuda(), "input must be a CUDA tensor");
+  AT_ASSERTM(rois.type().is_cuda(), "rois must be a CUDA tensor");
 
   auto num_rois = rois.size(0);
   auto channels = input.size(1);
@@ -162,8 +162,8 @@ at::Tensor ROIPool_backward_cuda(const at::Tensor& grad,
                                  const int channels,
                                  const int height,
                                  const int width) {
-  AT_ASSERT(grad.type().is_cuda(), "grad must be a CUDA tensor");
-  AT_ASSERT(rois.type().is_cuda(), "rois must be a CUDA tensor");
+  AT_ASSERTM(grad.type().is_cuda(), "grad must be a CUDA tensor");
+  AT_ASSERTM(rois.type().is_cuda(), "rois must be a CUDA tensor");
   // TODO add more checks
 
   auto num_rois = rois.size(0);