fix cuda bug in unit test

MadryLab · Nov 2, 2023 · 72d73a2 · 72d73a2
1 parent 8f30687
commit 72d73a2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_class.py b/tests/test_class.py
@@ -377,7 +377,7 @@ def test_custom_model_output(tmp_path, cpu_proj):
 def test_grad_wrt_last_layer(tmp_path):
     model = resnet18().cuda().eval()
     N = 5
-    batch = ch.randn(N, 3, 32, 32).cuda(), ch.randint(low=0, high=10, size=(N,)).cuda()
+    batch = ch.randn(N, 3, 32, 32).cuda(), ch.randint(low=0, high=10, size=(N,))
     traker = TRAKer(
         model=model,
         task="image_classification",