Check cuda script advances #26

pab1s · Jun 24, 2024 · 129f81e · 129f81e
1 parent f88a6e4
commit 129f81e
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/device_check.py b/device_check.py
@@ -0,0 +1,8 @@
+import torch
+
+def main():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+if __name__ == "__main__":
+    main()