Reduce ms_deformable_attn test memory usage (#1407)

open-mmlab · Oct 15, 2021 · 91b8478 · 91b8478
1 parent b484aba
commit 91b8478
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_ops/test_ms_deformable_attn.py b/tests/test_ops/test_ms_deformable_attn.py
@@ -151,7 +151,7 @@ def test_gradient_numerical(channels,
 
     N, M, _ = 1, 2, 2
     Lq, L, P = 2, 2, 2
-    shapes = torch.as_tensor([(6, 4), (3, 2)], dtype=torch.long).cuda()
+    shapes = torch.as_tensor([(3, 2), (2, 1)], dtype=torch.long).cuda()
     level_start_index = torch.cat((shapes.new_zeros(
         (1, )), shapes.prod(1).cumsum(0)[:-1]))
     S = sum([(H * W).item() for H, W in shapes])