Example doesn't run on gpu (when uncommenting)

[JamesOwers]

# dtype = torch.device("cuda:0") # Uncomment this to run on GPU
should be
# device = torch.device("cuda:0") # Uncomment this to run on GPU