Upload test_gpu.py
Browse files- test_gpu.py +5 -0
test_gpu.py
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
print('CUDA available:', torch.cuda.is_available())
|
| 3 |
+
if torch.cuda.is_available():
|
| 4 |
+
print('Device:', torch.cuda.get_device_name(0))
|
| 5 |
+
print('Memory:', torch.cuda.get_device_properties(0).total_memory / 1e9, 'GB')
|