Minor fix for cuda cli option in FFCV examples

ContinualAI · Oct 11, 2023 · 1ca45ad · 1ca45ad
1 parent a09c096
commit 1ca45ad
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/examples/ffcv/ffcv_enable.py b/examples/ffcv/ffcv_enable.py
@@ -24,7 +24,7 @@
 
 def main(cuda: int):
     # --- CONFIG
-    device = torch.device(f"cuda:{cuda}" if torch.cuda.is_available() else "cpu")
+    device = torch.device(f"cuda:{cuda}" if cuda >= 0 and torch.cuda.is_available() else "cpu")
     RNGManager.set_random_seeds(1234)
 
     benchmark_type = "cifar100"

diff --git a/examples/ffcv/ffcv_enable_rgb_compress.py b/examples/ffcv/ffcv_enable_rgb_compress.py
@@ -31,7 +31,7 @@
 
 
 def main(cuda: int):
-    device = torch.device(f"cuda:{cuda}" if torch.cuda.is_available() else "cpu")
+    device = torch.device(f"cuda:{cuda}" if cuda >= 0 and torch.cuda.is_available() else "cpu")
     RNGManager.set_random_seeds(1234)
 
     benchmark_type = "tinyimagenet"

diff --git a/examples/ffcv/ffcv_try_speed.py b/examples/ffcv/ffcv_try_speed.py
@@ -37,7 +37,7 @@
 
 def main(cuda: int):
     # --- CONFIG
-    device = torch.device(f"cuda:{cuda}" if torch.cuda.is_available() else "cpu")
+    device = torch.device(f"cuda:{cuda}" if cuda >= 0 and torch.cuda.is_available() else "cpu")
     RNGManager.set_random_seeds(1234)
 
     benchmark_type = "cifar100"