littlefig-benchmarks / test_gpu.py
ticketguy's picture
Upload test_gpu.py
cff6c87 verified
import torch
print('CUDA available:', torch.cuda.is_available())
if torch.cuda.is_available():
print('Device:', torch.cuda.get_device_name(0))
print('Memory:', torch.cuda.get_device_properties(0).total_memory / 1e9, 'GB')