remove device parameter in tests

jacobgil · Dec 18, 2023 · aa45c73 · aa45c73
1 parent e658038
commit aa45c73
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 6 deletions.
diff --git a/tests/test_context_release.py b/tests/test_context_release.py
@@ -57,8 +57,7 @@ def test_memory_usage_in_loop(numpy_image, batch_size, width, height,
     initial_memory = 0
     for i in range(100):
         with cam_method(model=model,
-                        target_layers=target_layers,
-                        device=None) as cam:
+                        target_layers=target_layers) as cam:
             grayscale_cam = cam(input_tensor=input_tensor,
                                 targets=targets,
                                 aug_smooth=aug_smooth,

diff --git a/tests/test_one_channel.py b/tests/test_one_channel.py
@@ -42,8 +42,7 @@ def test_memory_usage_in_loop(numpy_image, cam_method):
     print("input_tensor", input_tensor.shape)
     targets = None
     with cam_method(model=model,
-                    target_layers=target_layers,
-                    device=None) as cam:
+                    target_layers=target_layers) as cam:
         grayscale_cam = cam(input_tensor=input_tensor,
                             targets=targets)
         print(grayscale_cam.shape)
diff --git a/tests/test_run_all_models.py b/tests/test_run_all_models.py
@@ -64,8 +64,7 @@ def test_all_cam_models_can_run(numpy_image, batch_size, width, height,
         target_layers.append(eval(f"model.{layer}"))
 
     cam = cam_method(model=model,
-                     target_layers=target_layers,
-                     device=None)
+                     target_layers=target_layers)
     cam.batch_size = 4
     if target_category is None:
         targets = None