use_cuda = torch.cuda.is_available()
if use_cuda:
data, target = Variable(data.cuda(async=True, volatile=True)), Variable(target.cuda(async=True)) # On GPU
else:
data, target = Variable(data), Variable(target)
# You will get RuntimeError: expected CPU tensor (got CUDA tensor) if you dont do this