Warn if cuda present but not used

deepchecks · May 24, 2022 · 8be63cc · 8be63cc
1 parent 95d8481
commit 8be63cc
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/deepchecks/vision/base_checks.py b/deepchecks/vision/base_checks.py
@@ -42,7 +42,7 @@ def run(
         self,
         dataset: VisionData,
         model: Optional[nn.Module] = None,
-        device: Union[str, torch.device, None] = 'cpu',
+        device: Union[str, torch.device, None] = None,
         random_state: int = 42,
         n_samples: Optional[int] = 10_000
     ) -> CheckResult:
@@ -110,7 +110,7 @@ def run(
         train_dataset: VisionData,
         test_dataset: VisionData,
         model: Optional[nn.Module] = None,
-        device: Union[str, torch.device, None] = 'cpu',
+        device: Union[str, torch.device, None] = None,
         random_state: int = 42,
         n_samples: Optional[int] = 10_000
     ) -> CheckResult:
@@ -189,7 +189,7 @@ class ModelOnlyCheck(ModelOnlyBaseCheck):
     def run(
         self,
         model: nn.Module,
-        device: Union[str, torch.device, None] = 'cpu',
+        device: Union[str, torch.device, None] = None,
         random_state: int = 42
     ) -> CheckResult:
         """Run check."""

diff --git a/deepchecks/vision/context.py b/deepchecks/vision/context.py
@@ -59,7 +59,7 @@ def __init__(self,
                  model_name: str = '',
                  scorers: Mapping[str, Metric] = None,
                  scorers_per_class: Mapping[str, Metric] = None,
-                 device: Union[str, torch.device, None] = 'cpu',
+                 device: Union[str, torch.device, None] = None,
                  random_state: int = 42,
                  n_samples: int = None
                  ):
@@ -72,9 +72,14 @@ def __init__(self,
         if train and test:
             train.validate_shared_label(test)
 
+        if device is None:
+            device = 'cpu'
+            if torch.cuda.is_available():
+                warnings.warn('Checks will run on the cpu by default.'
+                              'To make use of cuda devices, use the device parameter in the run function.')
         self._device = torch.device(device) if isinstance(device, str) else (device if device else torch.device('cpu'))
-        self._prediction_formatter_error = {}
 
+        self._prediction_formatter_error = {}
         if model is not None:
             if not isinstance(model, nn.Module):
                 warnings.warn('Deepchecks can\'t validate that model is in evaluation state. Make sure it is to '