torch_gpu_test.py

torch_gpu_test.py

import torch
import sys
print('__Python VERSION:', sys.version)
print('__pyTorch VERSION:', torch.__version__)
print("Whether CUDA is available:", torch.cuda.is_available())
print('__CUDA VERSION', )
from subprocess import call
print('__CUDNN VERSION:', torch.backends.cudnn.version())
print('__Number CUDA Devices:', torch.cuda.device_count())
print('__Devices')
# call(["nvidia-smi", "--format=csv", "--query-gpu=index,name,driver_version,memory.total,memory.used,memory.free"])
print('Active CUDA Device: GPU', torch.cuda.current_device())
print('CUDA device')

if torch.cuda.is_available():
    device = torch.device('cuda')
var = torch.Tensor([1.0])
print(var.cuda())
print(var.to(device)) # <- preferred