Repository: deepchecks/deepchecks

Issue:
[FEAT] Add GPU instances for testing
**Is your feature request related to a problem? Please describe.**
We want to test the package on GPU instances. For that - we need a GPU instance on AWS that will be a github runner.
Then, we want all checks to work with gpu if gpu exists.

```
if torch.cuda.is_available():
  device = torch.device('cuda')
else:
  device = torch.device('cpu')
```




Discussion:

Resolving patch:
diff --git a/deepchecks/vision/metrics_utils/detection_precision_recall.py b/deepchecks/vision/metrics_utils/detection_precision_recall.py
--- a/deepchecks/vision/metrics_utils/detection_precision_recall.py
+++ b/deepchecks/vision/metrics_utils/detection_precision_recall.py
@@ -77,6 +77,11 @@ def update(self, output):
         y_pred, y = output
 
         for detected, ground_truth in zip(y_pred, y):
+            if isinstance(detected, torch.Tensor):
+                detected = detected.cpu()
+            if isinstance(ground_truth, torch.Tensor):
+                ground_truth = ground_truth.cpu()
+
             self._group_detections(detected, ground_truth)
             self.i += 1
 
diff --git a/deepchecks/vision/utils/transformations.py b/deepchecks/vision/utils/transformations.py
--- a/deepchecks/vision/utils/transformations.py
+++ b/deepchecks/vision/utils/transformations.py
@@ -139,8 +139,8 @@ def un_normalize_batch(tensor, mean: Sized, std: Sized, max_pixel_value: int = 2
     dim = len(mean)
     reshape_shape = (1, 1, 1, dim)
     max_pixel_value = [max_pixel_value] * dim
-    mean = torch.tensor(mean).reshape(reshape_shape)
-    std = torch.tensor(std).reshape(reshape_shape)
+    mean = torch.tensor(mean, device=tensor.device).reshape(reshape_shape)
+    std = torch.tensor(std, device=tensor.device).reshape(reshape_shape)
     tensor = (tensor * std) + mean
-    tensor = tensor * torch.tensor(max_pixel_value).reshape(reshape_shape)
+    tensor = tensor * torch.tensor(max_pixel_value, device=tensor.device).reshape(reshape_shape)
     return tensor.cpu().numpy()


Test patch:
diff --git a/tests/vision/base/test_context.py b/tests/vision/base/test_context.py
--- a/tests/vision/base/test_context.py
+++ b/tests/vision/base/test_context.py
@@ -33,14 +33,14 @@
 
 
 def test_vision_context_initialization_for_classification_task(mnist_dataset_train, mnist_dataset_test,
-                                                               trained_mnist):
+                                                               trained_mnist, device):
     # Act
     context = Context(
         train=mnist_dataset_train,
         test=mnist_dataset_test,
         model=trained_mnist,
         model_name='MNIST',
-        device='cpu',
+        device=device,
         prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter)
     )
 
@@ -52,20 +52,20 @@ def test_vision_context_initialization_for_classification_task(mnist_dataset_tra
         'model_name': equal_to('MNIST'),
         'device': all_of(
             instance_of(torch.device),
-            has_property('type', equal_to('cpu'))
+            has_property('type', equal_to(device.type))
         )
     }))
 
 
 def test_vision_context_initialization_for_object_detection_task(coco_train_visiondata, coco_test_visiondata,
-                                                                 trained_yolov5_object_detection):
+                                                                 trained_yolov5_object_detection, device):
     # Act
     context = Context(
         train=coco_train_visiondata,
         test=coco_test_visiondata,
         model=trained_yolov5_object_detection,
         model_name='COCO',
-        device='cpu',
+        device=device,
         prediction_formatter=DetectionPredictionFormatter(yolo_prediction_formatter)
     )
 
@@ -77,7 +77,7 @@ def test_vision_context_initialization_for_object_detection_task(coco_train_visi
         'model_name': equal_to('COCO'),
         'device': all_of(
             instance_of(torch.device),
-            has_property('type', equal_to('cpu'))
+            has_property('type', device.type)
         )
     }))
 
diff --git a/tests/vision/checks/distribution/heatmap_comparison_test.py b/tests/vision/checks/distribution/heatmap_comparison_test.py
--- a/tests/vision/checks/distribution/heatmap_comparison_test.py
+++ b/tests/vision/checks/distribution/heatmap_comparison_test.py
@@ -15,12 +15,12 @@
 from deepchecks.vision.checks.distribution import HeatmapComparison
 
 
-def test_object_detection(coco_train_visiondata, coco_test_visiondata):
+def test_object_detection(coco_train_visiondata, coco_test_visiondata, device):
     # Arrange
     check = HeatmapComparison()
 
     # Act
-    result = check.run(coco_train_visiondata, coco_test_visiondata)
+    result = check.run(coco_train_visiondata, coco_test_visiondata, device=device)
 
     # Assert
     brightness_diff = result.value["diff"]
@@ -32,12 +32,12 @@ def test_object_detection(coco_train_visiondata, coco_test_visiondata):
     assert_that(bbox_diff.max(), less_than(25))
 
 
-def test_classification(mnist_dataset_train, mnist_dataset_test):
+def test_classification(mnist_dataset_train, mnist_dataset_test, device):
     # Arrange
     check = HeatmapComparison()
 
     # Act
-    result = check.run(mnist_dataset_train, mnist_dataset_test)
+    result = check.run(mnist_dataset_train, mnist_dataset_test, device=device)
 
     # Assert
     brightness_diff = result.value["diff"]
@@ -45,12 +45,12 @@ def test_classification(mnist_dataset_train, mnist_dataset_test):
     assert_that(brightness_diff.max(), less_than(10))
 
 
-def test_object_detection_limit_classes(coco_train_visiondata, coco_test_visiondata):
+def test_object_detection_limit_classes(coco_train_visiondata, coco_test_visiondata, device):
     # Arrange
     check = HeatmapComparison(classes_to_display=['person'])
 
     # Act
-    result = check.run(coco_train_visiondata, coco_test_visiondata)
+    result = check.run(coco_train_visiondata, coco_test_visiondata, device=device)
 
     # Assert
     brightness_diff = result.value["diff"]
@@ -62,24 +62,24 @@ def test_object_detection_limit_classes(coco_train_visiondata, coco_test_visiond
     assert_that(bbox_diff.max(), less_than(22))
 
 
-def test_limit_classes_for_classification(mnist_dataset_train, mnist_dataset_test):
+def test_limit_classes_for_classification(mnist_dataset_train, mnist_dataset_test, device):
     # Arrange
     check = HeatmapComparison(classes_to_display=['0'])
 
     # Act & Assert
     assert_that(
-        calling(check.run).with_args(mnist_dataset_test, mnist_dataset_test),
+        calling(check.run).with_args(mnist_dataset_test, mnist_dataset_test, device=device),
         raises(DeepchecksNotSupportedError, 'Classes to display is only supported for object detection tasks.')
     )
 
 
-def test_limit_classes_nonexistant_class(coco_train_visiondata, coco_test_visiondata):
+def test_limit_classes_nonexistant_class(coco_train_visiondata, coco_test_visiondata, device):
     # Arrange
     check = HeatmapComparison(classes_to_display=['1000'])
 
     # Act & Assert
     assert_that(
-        calling(check.run).with_args(coco_train_visiondata, coco_test_visiondata),
+        calling(check.run).with_args(coco_train_visiondata, coco_test_visiondata, device=device),
         raises(DeepchecksValueError,
                r'Provided list of class ids to display \[\'1000\'\] not found in training dataset.')
     )
diff --git a/tests/vision/checks/distribution/image_dataset_drift_test.py b/tests/vision/checks/distribution/image_dataset_drift_test.py
--- a/tests/vision/checks/distribution/image_dataset_drift_test.py
+++ b/tests/vision/checks/distribution/image_dataset_drift_test.py
@@ -32,13 +32,13 @@ def pil_drift_formatter(batch):
     return [add_brightness(np.array(img)) for img in batch[0]]
 
 
-def test_no_drift_grayscale(mnist_dataset_train):
+def test_no_drift_grayscale(mnist_dataset_train, device):
     # Arrange
     train, test = mnist_dataset_train, mnist_dataset_train
     check = ImageDatasetDrift()
 
     # Act
-    result = check.run(train, test, random_state=42)
+    result = check.run(train, test, random_state=42, device=device)
 
     # Assert
     assert_that(result.value, has_entries({
@@ -55,13 +55,13 @@ def test_no_drift_grayscale(mnist_dataset_train):
     }))
 
 
-def test_drift_grayscale(mnist_dataset_train, mnist_dataset_test):
+def test_drift_grayscale(mnist_dataset_train, mnist_dataset_test, device):
     # Arrange
     train, test = mnist_dataset_train, mnist_dataset_test
     check = ImageDatasetDrift()
 
     # Act
-    result = check.run(train, test, random_state=42)
+    result = check.run(train, test, random_state=42, device=device)
 
     # Assert
     assert_that(result.value, has_entries({
@@ -78,7 +78,7 @@ def test_drift_grayscale(mnist_dataset_train, mnist_dataset_test):
     }))
 
 
-def test_no_drift_rgb(coco_train_dataloader, coco_test_dataloader):
+def test_no_drift_rgb(coco_train_dataloader, coco_test_dataloader, device):
     # Arrange
     train = VisionData(coco_train_dataloader, image_formatter=ImageFormatter(pil_formatter),
                        label_formatter=DetectionLabelFormatter())
@@ -88,7 +88,7 @@ def test_no_drift_rgb(coco_train_dataloader, coco_test_dataloader):
     check = ImageDatasetDrift()
 
     # Act
-    result = check.run(train, test, random_state=42)
+    result = check.run(train, test, random_state=42, device=device)
 
     # Assert
     assert_that(result.value, has_entries({
@@ -105,7 +105,7 @@ def test_no_drift_rgb(coco_train_dataloader, coco_test_dataloader):
     }))
 
 
-def test_with_drift_rgb(coco_train_dataloader, coco_test_dataloader):
+def test_with_drift_rgb(coco_train_dataloader, coco_test_dataloader, device):
     # Arrange
     train = VisionData(coco_train_dataloader, image_formatter=ImageFormatter(pil_drift_formatter),
                        label_formatter=DetectionLabelFormatter())
@@ -115,7 +115,7 @@ def test_with_drift_rgb(coco_train_dataloader, coco_test_dataloader):
     check = ImageDatasetDrift()
 
     # Act
-    result = check.run(train, test, random_state=42)
+    result = check.run(train, test, random_state=42, device=device)
 
     # Assert
     assert_that(result.value, has_entries({
diff --git a/tests/vision/checks/distribution/image_properties_drift_test.py b/tests/vision/checks/distribution/image_properties_drift_test.py
--- a/tests/vision/checks/distribution/image_properties_drift_test.py
+++ b/tests/vision/checks/distribution/image_properties_drift_test.py
@@ -33,11 +33,10 @@
 from deepchecks.vision.datasets.detection import coco
 
 
-
-def test_image_property_drift_check():
+def test_image_property_drift_check(device):
     train_dataset = coco.load_dataset(train=True, object_type='VisionData')
     test_dataset = coco.load_dataset(train=False, object_type='VisionData')
-    result = ImagePropertyDrift().run(train_dataset, test_dataset)
+    result = ImagePropertyDrift().run(train_dataset, test_dataset, device=device)
     assert_that(result, is_correct_image_property_drift_result())
 
 
@@ -55,14 +54,14 @@ def test_image_property_drift_initialization_with_list_of_unknown_image_properti
     )
 
 
-def test_image_property_drift_condition():
+def test_image_property_drift_condition(device):
     train_dataset = coco.load_dataset(train=True, object_type='VisionData')
     test_dataset = coco.load_dataset(train=False, object_type='VisionData')
 
     result = (
         ImagePropertyDrift()
         .add_condition_drift_score_not_greater_than()
-        .run(train_dataset, test_dataset)
+        .run(train_dataset, test_dataset, device=device)
     )
 
     assert_that(result, all_of(
diff --git a/tests/vision/checks/distribution/train_test_label_drift_test.py b/tests/vision/checks/distribution/train_test_label_drift_test.py
--- a/tests/vision/checks/distribution/train_test_label_drift_test.py
+++ b/tests/vision/checks/distribution/train_test_label_drift_test.py
@@ -15,13 +15,13 @@
 from deepchecks.vision.checks import TrainTestLabelDrift
 
 
-def test_no_drift_classification(mnist_dataset_train):
+def test_no_drift_classification(mnist_dataset_train, device):
     # Arrange
     train, test = mnist_dataset_train, mnist_dataset_train
     check = TrainTestLabelDrift()
 
     # Act
-    result = check.run(train, test)
+    result = check.run(train, test, device=device)
 
     # Assert
     assert_that(result.value, has_entries(
@@ -32,12 +32,12 @@ def test_no_drift_classification(mnist_dataset_train):
     ))
 
 
-def test_no_drift_object_detection(coco_train_visiondata):
+def test_no_drift_object_detection(coco_train_visiondata, device):
     # Arrange
     check = TrainTestLabelDrift()
 
     # Act
-    result = check.run(coco_train_visiondata, coco_train_visiondata)
+    result = check.run(coco_train_visiondata, coco_train_visiondata, device=device)
 
     # Assert
     assert_that(result.value, has_entries(
@@ -55,13 +55,13 @@ def test_no_drift_object_detection(coco_train_visiondata):
     ))
 
 
-def test_with_drift_classification(mnist_dataset_train, mnist_dataset_test):
+def test_with_drift_classification(mnist_dataset_train, mnist_dataset_test, device):
     # Arrange
     train, test = mnist_dataset_train, mnist_dataset_test
     check = TrainTestLabelDrift()
 
     # Act
-    result = check.run(train, test)
+    result = check.run(train, test, device=device)
 
     # Assert
     assert_that(result.value, has_entries(
@@ -73,12 +73,12 @@ def test_with_drift_classification(mnist_dataset_train, mnist_dataset_test):
     ))
 
 
-def test_with_drift_object_detection(coco_train_visiondata, coco_test_visiondata):
+def test_with_drift_object_detection(coco_train_visiondata, coco_test_visiondata, device):
     # Arrange
     check = TrainTestLabelDrift()
 
     # Act
-    result = check.run(coco_train_visiondata, coco_test_visiondata)
+    result = check.run(coco_train_visiondata, coco_test_visiondata, device=device)
 
     # Assert
     assert_that(result.value, has_entries(
@@ -96,12 +96,12 @@ def test_with_drift_object_detection(coco_train_visiondata, coco_test_visiondata
     ))
 
 
-def test_with_drift_object_detection_changed_num_samples(coco_train_visiondata, coco_test_visiondata):
+def test_with_drift_object_detection_changed_num_samples(coco_train_visiondata, coco_test_visiondata, device):
     # Arrange
     check = TrainTestLabelDrift(min_sample_size=32)
 
     # Act
-    result = check.run(coco_train_visiondata, coco_test_visiondata)
+    result = check.run(coco_train_visiondata, coco_test_visiondata, device=device)
 
     # Assert
     assert_that(result.value, has_entries(
@@ -119,12 +119,12 @@ def test_with_drift_object_detection_changed_num_samples(coco_train_visiondata,
     ))
 
 
-def test_with_drift_object_detection_changed_num_bins(coco_train_visiondata, coco_test_visiondata):
+def test_with_drift_object_detection_changed_num_bins(coco_train_visiondata, coco_test_visiondata, device):
     # Arrange
     check = TrainTestLabelDrift(default_num_bins=10)
 
     # Act
-    result = check.run(coco_train_visiondata, coco_test_visiondata)
+    result = check.run(coco_train_visiondata, coco_test_visiondata, device=device)
 
     # Assert
     assert_that(result.value, has_entries(
@@ -142,14 +142,14 @@ def test_with_drift_object_detection_changed_num_bins(coco_train_visiondata, coc
     ))
 
 
-def test_with_drift_object_detection_alternative_measurements(coco_train_visiondata, coco_test_visiondata):
+def test_with_drift_object_detection_alternative_measurements(coco_train_visiondata, coco_test_visiondata, device):
     # Arrange
     alternative_measurements = [
         {'name': 'test', 'method': lambda x, dataset: int(x[0][0]) if len(x) != 0 else 0, 'is_continuous': True}]
     check = TrainTestLabelDrift(alternative_label_measurements=alternative_measurements)
 
     # Act
-    result = check.run(coco_train_visiondata, coco_test_visiondata)
+    result = check.run(coco_train_visiondata, coco_test_visiondata, device=device)
 
     # Assert
     assert_that(result.value, has_entries(
diff --git a/tests/vision/checks/performance/class_performance_test.py b/tests/vision/checks/performance/class_performance_test.py
--- a/tests/vision/checks/performance/class_performance_test.py
+++ b/tests/vision/checks/performance/class_performance_test.py
@@ -19,13 +19,14 @@
 from hamcrest import assert_that, close_to, equal_to, is_in
 
 
-def test_mnist_largest(mnist_dataset_train, mnist_dataset_test, trained_mnist):
+def test_mnist_largest(mnist_dataset_train, mnist_dataset_test, trained_mnist, device):
     # Arrange
 
     check = ClassPerformance(n_to_show=2, show_only='largest')
     # Act
     result = check.run(mnist_dataset_train, mnist_dataset_test, trained_mnist,
-                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter))
+                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                       device=device)
     first_row = result.value.sort_values(by='Number of samples', ascending=False).iloc[0]
     # Assert
     assert_that(len(result.value), equal_to(8))
@@ -34,13 +35,14 @@ def test_mnist_largest(mnist_dataset_train, mnist_dataset_test, trained_mnist):
     assert_that(first_row['Class'], equal_to(1))
 
 
-def test_mnist_smallest(mnist_dataset_train, mnist_dataset_test, trained_mnist):
+def test_mnist_smallest(mnist_dataset_train, mnist_dataset_test, trained_mnist, device):
     # Arrange
 
     check = ClassPerformance(n_to_show=2, show_only='smallest')
     # Act
     result = check.run(mnist_dataset_train, mnist_dataset_test, trained_mnist,
-                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter))
+                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                       device=device)
     first_row = result.value.sort_values(by='Number of samples', ascending=True).iloc[0]
 
     # Assert
@@ -50,13 +52,14 @@ def test_mnist_smallest(mnist_dataset_train, mnist_dataset_test, trained_mnist):
     assert_that(first_row['Class'], equal_to(5))
 
 
-def test_mnist_worst(mnist_dataset_train, mnist_dataset_test, trained_mnist):
+def test_mnist_worst(mnist_dataset_train, mnist_dataset_test, trained_mnist, device):
     # Arrange
 
     check = ClassPerformance(n_to_show=2, show_only='worst')
     # Act
     result = check.run(mnist_dataset_train, mnist_dataset_test, trained_mnist,
-                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter))
+                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                       device=device)
     first_row = result.value.loc[result.value['Metric'] == 'Precision'].sort_values(by='Value', ascending=True).iloc[0]
 
     # Assert
@@ -64,13 +67,14 @@ def test_mnist_worst(mnist_dataset_train, mnist_dataset_test, trained_mnist):
     assert_that(first_row['Value'], close_to(0.977713, 0.05))
 
 
-def test_mnist_best(mnist_dataset_train, mnist_dataset_test, trained_mnist):
+def test_mnist_best(mnist_dataset_train, mnist_dataset_test, trained_mnist, device):
     # Arrange
 
     check = ClassPerformance(n_to_show=2, show_only='best')
     # Act
     result = check.run(mnist_dataset_train, mnist_dataset_test, trained_mnist,
-                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter))
+                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                       device=device)
     first_row = result.value.loc[result.value['Metric'] == 'Precision'].sort_values(by='Value', ascending=False).iloc[0]
 
     # Assert
@@ -78,13 +82,13 @@ def test_mnist_best(mnist_dataset_train, mnist_dataset_test, trained_mnist):
     assert_that(first_row['Value'], close_to(0.990854, 0.05))
 
 
-def test_coco_best(coco_train_visiondata, coco_test_visiondata, trained_yolov5_object_detection):
+def test_coco_best(coco_train_visiondata, coco_test_visiondata, trained_yolov5_object_detection, device):
     # Arrange
     pred_formatter = DetectionPredictionFormatter(yolo_prediction_formatter)
     check = ClassPerformance(n_to_show=2, show_only='best')
     # Act
     result = check.run(coco_train_visiondata, coco_test_visiondata,
-                       trained_yolov5_object_detection, prediction_formatter=pred_formatter)
+                       trained_yolov5_object_detection, prediction_formatter=pred_formatter, device=device)
 
     # Assert
     assert_that(len(result.value), equal_to(8))
@@ -100,13 +104,14 @@ def test_coco_best(coco_train_visiondata, coco_test_visiondata, trained_yolov5_o
     assert_that(first_row['Class'], is_in([29]))
 
 
-def test_class_list(mnist_dataset_train, mnist_dataset_test, trained_mnist):
+def test_class_list(mnist_dataset_train, mnist_dataset_test, trained_mnist, device):
     # Arrange
 
     check = ClassPerformance(class_list_to_show=[1])
     # Act
     result = check.run(mnist_dataset_train, mnist_dataset_test, trained_mnist,
-                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter))
+                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                       device=device)
 
     # Assert
     assert_that(len(result.value), equal_to(4))
diff --git a/tests/vision/checks/performance/confusion_matrix_test.py b/tests/vision/checks/performance/confusion_matrix_test.py
--- a/tests/vision/checks/performance/confusion_matrix_test.py
+++ b/tests/vision/checks/performance/confusion_matrix_test.py
@@ -18,24 +18,26 @@
 from deepchecks.vision.utils import ClassificationPredictionFormatter, DetectionPredictionFormatter
 
 
-def test_classification(mnist_dataset_train, trained_mnist):
+def test_classification(mnist_dataset_train, trained_mnist, device):
     # Arrange
     check = ConfusionMatrixReport()
     # Act
     result = check.run(mnist_dataset_train, trained_mnist,
-                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter))
+                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                       device=device)
     # Assert
     assert_that(result.value.shape, equal_to((10, 10)))
 
 
-def test_detection(coco_train_visiondata, trained_yolov5_object_detection):
+def test_detection(coco_train_visiondata, trained_yolov5_object_detection, device):
     # Arrange
     pred_formatter = DetectionPredictionFormatter(yolo_prediction_formatter)
     check = ConfusionMatrixReport()
     # Act
     result = check.run(coco_train_visiondata,
                        trained_yolov5_object_detection,
-                       prediction_formatter=pred_formatter)
+                       prediction_formatter=pred_formatter,
+                       device=device)
 
     # Assert
     assert_that(result.value.shape, equal_to((81, 81)))
diff --git a/tests/vision/checks/performance/cv_simple_model_comparison_test.py b/tests/vision/checks/performance/cv_simple_model_comparison_test.py
--- a/tests/vision/checks/performance/cv_simple_model_comparison_test.py
+++ b/tests/vision/checks/performance/cv_simple_model_comparison_test.py
@@ -15,13 +15,14 @@
 from hamcrest import assert_that, close_to, equal_to, is_in, calling, raises
 
 
-def test_mnist_prior_strategy(mnist_dataset_train, mnist_dataset_test, trained_mnist):
+def test_mnist_prior_strategy(mnist_dataset_train, mnist_dataset_test, trained_mnist, device):
     # Arrange
 
     check = SimpleModelComparison(strategy='prior', n_to_show=2, show_only='largest')
     # Act
     result = check.run(mnist_dataset_train, mnist_dataset_test, trained_mnist,
-                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter))
+                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                       device=device)
     first_row = result.value.sort_values(by='Number of samples', ascending=False).iloc[0]
     # Assert
     assert_that(len(result.value), equal_to(8))
@@ -39,12 +40,13 @@ def test_mnist_not_exist_strategy(mnist_dataset_train, mnist_dataset_test, train
     )
 
 
-def test_mnist_most_frequent(mnist_dataset_train, mnist_dataset_test, trained_mnist):
+def test_mnist_most_frequent(mnist_dataset_train, mnist_dataset_test, trained_mnist, device):
     # Arrange
     check = SimpleModelComparison(strategy='most_frequent', n_to_show=2, show_only='largest')
     # Act
     result = check.run(mnist_dataset_train, mnist_dataset_test, trained_mnist,
-                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter))
+                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                       device=device)
     first_row = result.value.sort_values(by='Number of samples', ascending=False).iloc[0]
     # Assert
     assert_that(len(result.value), equal_to(8))
@@ -53,21 +55,23 @@ def test_mnist_most_frequent(mnist_dataset_train, mnist_dataset_test, trained_mn
     assert_that(first_row['Class'], equal_to(1))
 
 
-def test_mnist_uniform(mnist_dataset_train, mnist_dataset_test, trained_mnist):
+def test_mnist_uniform(mnist_dataset_train, mnist_dataset_test, trained_mnist, device):
     # Arrange
     check = SimpleModelComparison(strategy='uniform', n_to_show=2, show_only='largest')
     # Act
     result = check.run(mnist_dataset_train, mnist_dataset_test, trained_mnist,
-                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter))
+                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                       device=device)
     # Assert
     assert_that(len(result.value), equal_to(8))
 
 
-def test_mnist_stratified(mnist_dataset_train, mnist_dataset_test, trained_mnist):
+def test_mnist_stratified(mnist_dataset_train, mnist_dataset_test, trained_mnist, device):
     # Arrange
     check = SimpleModelComparison(strategy='stratified', n_to_show=2, show_only='largest')
     # Act
     result = check.run(mnist_dataset_train, mnist_dataset_test, trained_mnist,
-                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter))
+                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                       device=device)
     # Assert
     assert_that(len(result.value), equal_to(8))
\ No newline at end of file
diff --git a/tests/vision/checks/performance/mean_average_precision_report_test.py b/tests/vision/checks/performance/mean_average_precision_report_test.py
--- a/tests/vision/checks/performance/mean_average_precision_report_test.py
+++ b/tests/vision/checks/performance/mean_average_precision_report_test.py
@@ -18,19 +18,20 @@
 from deepchecks.vision.utils import DetectionPredictionFormatter, ClassificationPredictionFormatter
 
 
-def test_mnist_error(mnist_dataset_test, trained_mnist):
+def test_mnist_error(mnist_dataset_test, trained_mnist, device):
     # Arrange
     check = MeanAveragePrecisionReport()
     # Act
     assert_that(
         calling(check.run
                 ).with_args(mnist_dataset_test, trained_mnist,
-                            prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter)),
+                            prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                            device=device),
         raises(ModelValidationError, r'Check is irrelevant for task of type TaskType.CLASSIFICATION')
     )
 
 
-def test_coco(coco_test_visiondata, trained_yolov5_object_detection):
+def test_coco(coco_test_visiondata, trained_yolov5_object_detection, device):
     # Arrange
     pred_formatter = DetectionPredictionFormatter(yolo_prediction_formatter)
     check = MeanAveragePrecisionReport() \
@@ -39,7 +40,7 @@ def test_coco(coco_test_visiondata, trained_yolov5_object_detection):
 
     # Act
     result = check.run(coco_test_visiondata,
-                       trained_yolov5_object_detection, prediction_formatter=pred_formatter)
+                       trained_yolov5_object_detection, prediction_formatter=pred_formatter, device=device)
 
     # Assert
     df = result.value
@@ -74,14 +75,15 @@ def test_coco(coco_test_visiondata, trained_yolov5_object_detection):
     ))
 
 
-def test_coco_area_param(coco_test_visiondata, trained_yolov5_object_detection):
+def test_coco_area_param(coco_test_visiondata, trained_yolov5_object_detection, device):
     # Arrange
     pred_formatter = DetectionPredictionFormatter(yolo_prediction_formatter)
     check = MeanAveragePrecisionReport(area_range=(40**2, 100**2))
 
     # Act
     result = check.run(coco_test_visiondata,
-                       trained_yolov5_object_detection, prediction_formatter=pred_formatter)
+                       trained_yolov5_object_detection, prediction_formatter=pred_formatter,
+                       device=device)
 
     # Assert
     df = result.value
diff --git a/tests/vision/checks/performance/mean_average_recall_report_test.py b/tests/vision/checks/performance/mean_average_recall_report_test.py
--- a/tests/vision/checks/performance/mean_average_recall_report_test.py
+++ b/tests/vision/checks/performance/mean_average_recall_report_test.py
@@ -18,19 +18,20 @@
 from deepchecks.vision.utils import DetectionPredictionFormatter, ClassificationPredictionFormatter
 
 
-def test_mnist_error(mnist_dataset_test, trained_mnist):
+def test_mnist_error(mnist_dataset_test, trained_mnist, device):
     # Arrange
     check = MeanAverageRecallReport()
     # Act
     assert_that(
         calling(check.run
                 ).with_args(mnist_dataset_test, trained_mnist,
-                            prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter)),
+                            prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                            device=device),
             raises(ModelValidationError, r'Check is irrelevant for task of type TaskType.CLASSIFICATION')
     )
 
 
-def test_coco(coco_test_visiondata, trained_yolov5_object_detection):
+def test_coco(coco_test_visiondata, trained_yolov5_object_detection, device):
     # Arrange
     pred_formatter = DetectionPredictionFormatter(yolo_prediction_formatter)
     check = MeanAverageRecallReport() \
@@ -39,7 +40,7 @@ def test_coco(coco_test_visiondata, trained_yolov5_object_detection):
 
     # Act
     result = check.run(coco_test_visiondata,
-                       trained_yolov5_object_detection, prediction_formatter=pred_formatter)
+                       trained_yolov5_object_detection, prediction_formatter=pred_formatter, device=device)
 
     # Assert
     df = result.value
@@ -74,14 +75,14 @@ def test_coco(coco_test_visiondata, trained_yolov5_object_detection):
     ))
 
 
-def test_coco_area_param(coco_test_visiondata, trained_yolov5_object_detection):
+def test_coco_area_param(coco_test_visiondata, trained_yolov5_object_detection, device):
     # Arrange
     pred_formatter = DetectionPredictionFormatter(yolo_prediction_formatter)
     check = MeanAverageRecallReport(area_range=(40**2, 100**2))
 
     # Act
     result = check.run(coco_test_visiondata,
-                       trained_yolov5_object_detection, prediction_formatter=pred_formatter)
+                       trained_yolov5_object_detection, prediction_formatter=pred_formatter, device=device)
 
     # Assert
     df = result.value
diff --git a/tests/vision/checks/performance/robustness_report_test.py b/tests/vision/checks/performance/robustness_report_test.py
--- a/tests/vision/checks/performance/robustness_report_test.py
+++ b/tests/vision/checks/performance/robustness_report_test.py
@@ -28,7 +28,7 @@
 from tests.vision.vision_conftest import *
 
 
-def test_mnist(mnist_dataset_train, trained_mnist):
+def test_mnist(mnist_dataset_train, trained_mnist, device):
     # Arrange
     # Create augmentations without randomness to get fixed metrics results
     augmentations = [
@@ -38,7 +38,8 @@ def test_mnist(mnist_dataset_train, trained_mnist):
     check = RobustnessReport(augmentations=augmentations)
     # Act
     result = check.run(mnist_dataset_train, trained_mnist,
-                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter))
+                       prediction_formatter=ClassificationPredictionFormatter(mnist_prediction_formatter),
+                       device=device)
     # Assert
     assert_that(result.value, has_entries({
         'RandomBrightnessContrast': has_entries({
@@ -52,7 +53,7 @@ def test_mnist(mnist_dataset_train, trained_mnist):
     }))
 
 
-def test_coco_and_condition(coco_train_visiondata, trained_yolov5_object_detection):
+def test_coco_and_condition(coco_train_visiondata, trained_yolov5_object_detection, device):
     """Because of the large running time, instead of checking the conditions in separated tests, combining a few
     tests into one."""
     # Arrange
@@ -67,7 +68,8 @@ def test_coco_and_condition(coco_train_visiondata, trained_yolov5_object_detecti
     check.add_condition_degradation_not_greater_than(0.01)
 
     # Act
-    result = check.run(coco_train_visiondata, trained_yolov5_object_detection, prediction_formatter=pred_formatter)
+    result = check.run(coco_train_visiondata, trained_yolov5_object_detection, prediction_formatter=pred_formatter,
+                       device=device)
     # Assert
     assert_that(result.value, has_entries({
         'HueSaturationValue': has_entries({
@@ -88,7 +90,7 @@ def test_coco_and_condition(coco_train_visiondata, trained_yolov5_object_detecti
     ))
 
 
-def test_dataset_not_augmenting_labels(coco_train_visiondata, trained_yolov5_object_detection):
+def test_dataset_not_augmenting_labels(coco_train_visiondata, trained_yolov5_object_detection, device):
     # Arrange
     vision_data = coco_train_visiondata.copy()
     dataset = vision_data.get_data_loader().dataset
@@ -107,11 +109,12 @@ def new_apply(self, img, bboxes):
     msg = r'Found that labels have not been affected by adding augmentation to field "transforms". This might be ' \
           r'a problem with the implementation of `Dataset.__getitem__`. label value: .*'
     assert_that(calling(check.run).with_args(vision_data, trained_yolov5_object_detection,
-                                             prediction_formatter=pred_formatter),
+                                             prediction_formatter=pred_formatter,
+                                             device=device),
                 raises(DeepchecksValueError, msg))
 
 
-def test_dataset_not_augmenting_data(coco_train_visiondata, trained_yolov5_object_detection):
+def test_dataset_not_augmenting_data(coco_train_visiondata, trained_yolov5_object_detection, device):
     # Arrange
     vision_data = coco_train_visiondata.copy()
     dataset = vision_data.get_data_loader().dataset
@@ -127,5 +130,6 @@ def new_apply(self, img, bboxes):
     msg = r'Found that images have not been affected by adding augmentation to field "transforms". This might be a ' \
           r'problem with the implementation of Dataset.__getitem__'
     assert_that(calling(check.run).with_args(vision_data, trained_yolov5_object_detection,
-                                             prediction_formatter=pred_formatter),
+                                             prediction_formatter=pred_formatter,
+                                             device=device),
                 raises(DeepchecksValueError, msg))
diff --git a/tests/vision/utils_tests/metrics_test.py b/tests/vision/utils_tests/metrics_test.py
--- a/tests/vision/utils_tests/metrics_test.py
+++ b/tests/vision/utils_tests/metrics_test.py
@@ -8,7 +8,6 @@
 # along with Deepchecks.  If not, see <http://www.gnu.org/licenses/>.
 # ----------------------------------------------------------------------------
 #
-import torch
 from hamcrest import has_items, assert_that, has_length, close_to
 
 from deepchecks.vision.datasets.detection.coco import yolo_prediction_formatter
@@ -18,31 +17,31 @@
 from deepchecks.vision import VisionData
 
 
-def test_default_ap_ignite_complient(coco_test_visiondata: VisionData, trained_yolov5_object_detection):
+def test_default_ap_ignite_complient(coco_test_visiondata: VisionData, trained_yolov5_object_detection, device):
     res = calculate_metrics({'AveragePrecision': AveragePrecision()},
                             coco_test_visiondata, trained_yolov5_object_detection,
                             prediction_formatter=DetectionPredictionFormatter(yolo_prediction_formatter),
-                            device=torch.device('cpu'))
+                            device=device)
     assert_that(res.keys(), has_length(1))
     assert_that(res['AveragePrecision'], has_length(59))
 
 
-def test_ar_ignite_complient(coco_test_visiondata: VisionData, trained_yolov5_object_detection):
+def test_ar_ignite_complient(coco_test_visiondata: VisionData, trained_yolov5_object_detection, device):
     res = calculate_metrics({'AveragePrecision': AveragePrecision(return_option=1)},
                             coco_test_visiondata, trained_yolov5_object_detection,
                             prediction_formatter=DetectionPredictionFormatter(yolo_prediction_formatter),
-                            device=torch.device('cpu'))
+                            device=device)
 
     assert_that(res.keys(), has_length(1))
     assert_that(res['AveragePrecision'], has_length(59))
 
 
-def test_equal_pycocotools(coco_test_visiondata: VisionData, trained_yolov5_object_detection):
+def test_equal_pycocotools(coco_test_visiondata: VisionData, trained_yolov5_object_detection, device):
     metric = AveragePrecision(return_option=None)
     for batch in coco_test_visiondata.get_data_loader():
         label = coco_test_visiondata.label_formatter(batch)
         prediction = DetectionPredictionFormatter(yolo_prediction_formatter)(batch, trained_yolov5_object_detection,
-                                                                             torch.device('cpu'))
+                                                                             device)
         metric.update((prediction, label))
     res = metric.compute()[0]
 
diff --git a/tests/vision/vision_conftest.py b/tests/vision/vision_conftest.py
--- a/tests/vision/vision_conftest.py
+++ b/tests/vision/vision_conftest.py
@@ -29,7 +29,8 @@
 torch.hub.set_dir(str(PROJECT_DIR))
 
 
-__all__ = ['mnist_data_loader_train',
+__all__ = ['device',
+           'mnist_data_loader_train',
            'mnist_dataset_train',
            'mnist_data_loader_test',
            'mnist_dataset_train_imgaug',
@@ -45,6 +46,16 @@
         ]
 
 
+@pytest.fixture(scope='session')
+def device():
+    if torch.cuda.is_available():
+        device = torch.device('cuda:0')  # pylint: disable=redefined-outer-name
+    else:
+        device = torch.device('cpu')  # pylint: disable=redefined-outer-name
+
+    return device
+
+
 @pytest.fixture(scope='session')
 def mnist_data_loader_train():
     return load_mnist_dataset(train=True, object_type='DataLoader')
@@ -81,8 +92,8 @@ def mnist_dataset_train_imgaug():
 
 
 @pytest.fixture(scope='session')
-def trained_yolov5_object_detection():
-    return load_yolov5_model()
+def trained_yolov5_object_detection(device):  # pylint: disable=redefined-outer-name
+    return load_yolov5_model(device=device)
 
 
 @pytest.fixture(scope='session')