Testing for GPU availability

torch_gpu_test.py

import torch

device = torch.device('cuda:0') if torch.cuda.is_available() else 'cpu'

print(f'Using: {device}')
print('GPU Memory Usage:')
print('Allocated:', round(torch.cuda.memory_allocated(0) / 1024**3, 1), 'GB')
print('Cached:   ', round(torch.cuda.memory_reserved(0) / 1024**3, 1), 'GB')

# Run a trivial task
x = torch.rand(10).to(device)
x = x.sum().item()
print(f'Sum random numbers from {device}:', x)