Create device/get_device_id.py

device/get_device_id.py

@@ -0,0 +1,7 @@
+import torch
+
+def get_device_id():
+    if torch.cuda.is_available():
+        return 0
+
+    return -1