Skip to content

Instantly share code, notes, and snippets.

@jgoodie

jgoodie/set_cuda_device_check_gpu.py

Created May 25, 2024 23:34
Show Gist options
  • Save jgoodie/d6d434b814e17f14c262f999de2d382d to your computer and use it in GitHub Desktop.
Save jgoodie/d6d434b814e17f14c262f999de2d382d to your computer and use it in GitHub Desktop.
Download ZIP
Raw
set_cuda_device_check_gpu.py
# setting device on GPU if available, else CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
print()
#Additional Info when using cuda
dev = 0
if device.type == 'cuda':
print(torch.cuda.get_device_name(dev))
print('Memory Usage:')
print('Allocated:', round(torch.cuda.memory_allocated(dev)/1024**3,1), 'GB')
print('Cached: ', round(torch.cuda.memory_reserved(dev)/1024**3,1), 'GB')
print("|=======================================|")
#Additional Info when using cuda
dev = 1
if device.type == 'cuda':
print(torch.cuda.get_device_name(dev))
print('Memory Usage:')
print('Allocated:', round(torch.cuda.memory_allocated(dev)/1024**3,1), 'GB')
print('Cached: ', round(torch.cuda.memory_reserved(dev)/1024**3,1), 'GB')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment