make sure current device is set when no gpus availabble

neuralmagic · Jul 24, 2023 · c0cdb45 · c0cdb45
1 parent 564e5ae
commit c0cdb45
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/deepsparse/yolov8/utils/validation/deepsparse_validator.py b/src/deepsparse/yolov8/utils/validation/deepsparse_validator.py
@@ -7,6 +7,7 @@
 
 from tqdm import tqdm
 
+import torch
 from deepsparse.yolov8.utils.validation.helpers import schema_to_tensor
 from ultralytics.yolo.data.utils import check_det_dataset
 from ultralytics.yolo.utils import LOGGER, TQDM_BAR_FORMAT, callbacks
@@ -34,6 +35,8 @@ def __call__(self, classes: Dict[int, str]):
         # for validation when self.training is True
         callbacks.add_integration_callbacks(self)
         self.run_callbacks("on_val_start")
+        if not torch.cuda.is_available():
+            self.args.device = "cpu"
         self.device = select_device(self.args.device, self.args.batch)
         self.data = check_det_dataset(self.args.data)
         if isinstance(self.data["path"], str):