MAINT: only call cuda current device if there is GPU

Lycodeboy · Nov 3, 2020 · 1d6a3cb · 1d6a3cb
1 parent c77acd6
commit 1d6a3cb
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/utils.py b/utils.py
@@ -8,7 +8,10 @@
 import random
 
 import torch
-torch.cuda.current_device()
+
+if torch.cuda.is_available():
+    torch.cuda.current_device()
+
 import torchvision.transforms.functional as TF
 from torch.utils.data import Dataset, DataLoader, Subset
 from torchvision import transforms, utils