fix cuda bug in unit test part 3 oops

MadryLab · Nov 2, 2023 · d9b292c · d9b292c
1 parent be57466
commit d9b292c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_class.py b/tests/test_class.py
@@ -375,7 +375,7 @@ def test_custom_model_output(tmp_path, cpu_proj):
 
 
 def test_grad_wrt_last_layer(tmp_path):
-    model = resnet18().cuda().eval()
+    model = resnet18().eval()
     N = 5
     batch = ch.randn(N, 3, 32, 32).cuda(), ch.randint(low=0, high=10, size=(N,))
     traker = TRAKer(
@@ -400,7 +400,7 @@ def test_grad_wrt_last_layer(tmp_path):
 
 @pytest.mark.cuda
 def test_grad_wrt_last_layer_cuda(tmp_path):
-    model = resnet18().eval()
+    model = resnet18().cuda().eval()
     N = 5
     batch = ch.randn(N, 3, 32, 32).cuda(), ch.randint(low=0, high=10, size=(N,)).cuda()
     traker = TRAKer(