import torch
torch.cuda.is_available()

print('cuda available: ', torch.cuda.is_available())
print('device count: ', torch.cuda.device_count())
print('cuda current device: ', torch.cuda.current_device())
print('cude device name: ', torch.cuda.get_device_name(0))
print('allocated memory: ', torch.cuda.memory_allocated())
print('allocated memory: ', torch.cuda.memory_cached())

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