[fix] Add the assertion of batch_size in tin_shift.py

open-mmlab · Feb 11, 2022 · 363c819 · 363c819
1 parent b8ab4fd
commit 363c819
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/mmcv/ops/tin_shift.py b/mmcv/ops/tin_shift.py
@@ -18,6 +18,11 @@ class TINShiftFunction(Function):
 
     @staticmethod
     def forward(ctx, input, shift):
+        batches_input = input.size(0)
+        batches_shift = shift.size(0)
+        if batches_input != batches_shift:
+            raise ValueError('batches_input should be equal to batches_shift, '
+                             f'but got batches_input={batches_input} and batches_shift={batches_shift}.')
         C = input.size(2)
         num_segments = shift.size(1)
         if C // num_segments <= 0 or C % num_segments != 0:

diff --git a/tests/test_ops/test_tin_shift.py b/tests/test_ops/test_tin_shift.py
@@ -184,7 +184,7 @@ def _test_tinshift_assert(device, dtype):
     except ModuleNotFoundError:
         pytest.skip('TINShift op is not successfully compiled')
 
-    inputs = [torch.rand(2, 3, 4, 2), torch.rand(2, 3, 4, 2)]
+    inputs = [torch.rand(2, 3, 4, 2), torch.rand(2, 3, 4, 2), torch.rand(1, 3, 4, 2)]
     shifts = [torch.rand(2, 3), torch.rand(2, 5)]
 
     for x, shift in zip(inputs, shifts):