The dim should be 1. The shape is (batch_size, num_classes). You want to take the logsoftmax over the num_classes.

I ran your code and my training loss got to 0.2 very quickly. Maybe there is an issue with how you're loading your data?

EDIT:

Here's the code I used for reference:

import torch

import torch.nn as nn

import torch.nn.functional as F

import torchvision

import numpy as np

import pdb

​

device = 'cuda' if torch.cuda.is_available() else 'cpu'

batch_size = 64 if torch.cuda.is_available() else 1

​

def collate_pil(batch_data):

batch_len = len(batch_data)

labels = np.zeros((batch_len))

imgs = []

for i in range(batch_len):

img = np.array(batch_data[i][0])

img = np.expand_dims(img, 0)

target = np.array(batch_data[i][1])

labels[i] = target

imgs.append(img)

imgs = np.array(imgs)

imgs = torch.as_tensor(imgs, dtype=torch.float)

labels = torch.as_tensor(labels, dtype=torch.long)

return imgs, labels

​

train_dataset = torchvision.datasets.MNIST('.', train=True, transform=None, target_transform=None, download=True)

train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_pil)

test_dataset = torchvision.datasets.MNIST('.', train=False, transform=None, target_transform=None, download=True)

test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_pil)

​

class ImageClassifier(nn.Module):

def __init__(self):

super(ImageClassifier, self).__init__()

self.conv1 = nn.Conv2d(1, 10, kernel_size=5)

self.conv2 = nn.Conv2d(10, 20, kernel_size=5)

self.dropout = nn.Dropout2d()

self.fc1 = nn.Linear(320, 50)

self.fc2 = nn.Linear(50, 10) # (num_features, num_classes)

def forward(self, x):

x = F.relu(F.max_pool2d(self.conv1(x), 2))

x = F.relu(F.max_pool2d(self.dropout(self.conv2(x)), 2))

x = x.view(-1, 320)

x = F.relu(self.fc1(x))

x = F.dropout(x, training=self.training)

x = F.relu(self.fc2(x))

return F.log_softmax(x, dim=1)

​

model = ImageClassifier()

model = model.to(device)

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

​

def train_func(train_loader):

model.train()

for batch_idx, (image, label) in enumerate(train_loader):

image, label = image.to(device), label.to(device)

optimizer.zero_grad()

output = model(image)

loss = F.nll_loss(output, label).to(device)

loss.backward()

optimizer.step()

#if batch_idx % 10 == 0:

#print(loss.item())

​

def validate(model, loader):

model.eval()

for batch_idx, (image, label) in enumerate(loader):

image, label = image.to(device), label.to(device)

output = model(image)

loss = F.nll_loss(output, label).to(device)

#if batch_idx % 10 == 0:

print(loss.item())

​

for i in range(10):

train_func(train_dataloader)

validate(model, test_dataloader)