Use cuda if available, otherwise use cpu

mauna-ai · Mar 22, 2021 · cfd8424 · cfd8424
1 parent 82191b9
commit cfd8424
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/comet/data/config.py b/comet/data/config.py
@@ -1,7 +1,8 @@
 import json
+import torch
 from comet.utils import DD
 
-device = "cuda"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 save = False
 test_save = False