Test if pytorch has cuda enabled.

torch_test_cuda.py

# Simple test to check if pytorch is installed and can create matrices

import torch
x = torch.rand(5, 3)
print(x)
print("\n")

if torch.cuda.is_available():
    print("Is PyTorch CUDA installed: " + str(torch.cuda.is_available()))
    print("Number of Cuda Devices: " + str(torch.cuda.device_count()))
    print("Device Name: "+ str(torch.cuda.get_device_name(0)))
else:
    print("Is PyTorch CUDA installed: " + str(torch.cuda.is_available()))
    
    
# Creating a sample tensor
x = torch.randint(1, 1000, (100, 100))
# Checking the device name: will return ‘CPU’ by default
print("Device Name: " , x.device)
# Applying tensor operation
res_cpu = x ** 2
# Transferring tensor to GPU
x = x.to(torch.device('cuda'))
# Checking the device name: will return ‘cuda:0’
print("Device Name after transferring: ", x.device)
# Applying same tensor operation
res_gpu = x ** 2
# Transferring tensor from GPU to CPU
x.cpu()


# 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
if device.type == 'cuda':
    print(torch.cuda.get_device_name(0))
    print('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')