Here is the network configuration
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net().to(device)

And here is train code
for epoch in range(2):  # loop over the dataset multiple times

    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # get the inputs
        inputs, labels = data
        inputs = inputs.to(device)
        labels = labels.to(device)

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        # print('passed')
        optimizer.step()


        # print statistics
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

print('Finished Training')

And here is the error message
RuntimeError                              Traceback (most recent call last)
<ipython-input-20-6350afc1b67e> in <module>()
     17         loss.backward()
     18         print('pass')
---> 19         optimizer.step()
     20 
     21 

/usr/local/lib/python3.5/dist-packages/torch/optim/sgd.py in step(self, closure)
     99                     else:
    100                         buf = param_state['momentum_buffer']
--> 101                         buf.mul_(momentum).add_(1 - dampening, d_p)
    102                     if nesterov:
    103                         d_p = d_p.add(momentum, buf)

RuntimeError: Expected object of type torch.FloatTensor but found type torch.cuda.FloatTensor for argument #4 'other'

I’m new to PyTorch, and I just can’t figure out how can this still want torch.FloatTensor instead of torch.cuda.FloatTensor after I called .cuda().
Thanks.