runnning on cuda:0

  0%|                                                  | 0/4000 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "mnist_classifier.py", line 58, in <module>
    output = net(X.view(-1, 784))  # pass in the reshaped batch (recall they are 28x28 atm)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "mnist_classifier.py", line 34, in forward
    x = F.relu(self.fc1(x))
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/linear.py", line 87, in forward
    return F.linear(input, self.weight, self.bias)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py", line 1370, in linear
    ret = torch.addmm(bias, input, weight.t())
RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'mat1' in call to _th_addmm