[Fix] Fix the unit test of correlation op (open-mmlab#1659)

joemathai · Jan 13, 2022 · 227037f · 227037f
1 parent ccdc61c
commit 227037f
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/tests/test_ops/test_correlation.py b/tests/test_ops/test_correlation.py
@@ -30,10 +30,13 @@ def _test_correlation(self, dtype=torch.float):
         out = layer(input1, input2)
         out.backward(torch.ones_like(out))
 
-        gt_out = torch.tensor(_gt_out, dtype=dtype)
-        assert_equal_tensor(out.cpu(), gt_out)
-        assert_equal_tensor(input1.grad.detach().cpu(), input2.cpu())
-        assert_equal_tensor(input2.grad.detach().cpu(), input1.cpu())
+        # `eq_cpu` is not implemented for 'Half' in torch1.5.0,
+        # so we need to make a comparison for cuda tensor
+        # rather than cpu tensor
+        gt_out = torch.tensor(_gt_out, dtype=dtype).cuda()
+        assert_equal_tensor(out, gt_out)
+        assert_equal_tensor(input1.grad.detach(), input2)
+        assert_equal_tensor(input2.grad.detach(), input1)
 
     @pytest.mark.skipif(
         not torch.cuda.is_available(), reason='requires CUDA support')