Validations add call to cuda.init()

pytorch · Jul 24, 2024 · 2538602 · 2538602
1 parent c6fcf2d
commit 2538602
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/test/smoke_test/smoke_test.py b/test/smoke_test/smoke_test.py
@@ -177,6 +177,12 @@ def smoke_test_cuda(package: str, runtime_error_check: str, torch_compile_check:
         print(f"torch cudnn: {torch.backends.cudnn.version()}")
         print(f"cuDNN enabled? {torch.backends.cudnn.enabled}")
 
+        torch.cuda.init()
+        print(f"CUDA initialized successfully")
+        print(f"Number of CUDA devices: {torch.cuda.device_count()}")
+        for i in range(torch.cuda.device_count()):
+            print(f"Device {i}: {torch.cuda.get_device_name(i)}")
+
         # nccl is availbale only on Linux
         if (sys.platform in ["linux", "linux2"]):
             print(f"torch nccl version: {torch.cuda.nccl.version()}")