Pytorch MNIST code is returning IndexError
I've followed the Pytorch documentation and have made an extremely simple classifier for the MNIST dataset. Below is my code:
import numpy as np

import torch
import torchvision
from torchvision import transforms, datasets

import torch.nn as nn
import torch.nn.functional as F

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
    ])

train = datasets.MNIST('', train=True, download=True, transform=transform)
test = datasets.MNIST('', train=False, download=True, transform=transform)
trainset = torch.utils.data.DataLoader(train, batch_size=1, shuffle=True)
testset = torch.utils.data.DataLoader(test, batch_size=1, shuffle=False)

class Classifier(nn.Module):
    def __init__(self, D_in, H, D_out):
        super(Classifier, self).__init__()
        self.linear_1 = torch.nn.Linear(D_in, H)
        self.linear_2 = torch.nn.Linear(H, D_out)

    def forward(self, x):
        x = self.linear_1(x).clamp(min=0)
        x = self.linear_2(x)
        return F.log_softmax(x, dim=1)


net = Classifier(28*28, 128, 10)
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)

for epoch in range(3):
    running_loss = 0.0
    for X, label in iter(trainset):
        X = X.view(28*28, -1)

        optimizer.zero_grad()

        output = net(torch.flatten(X))
        loss = nn.CrossEntropyLoss(output, label)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print(f'[{epoch + 1}, {i + 1}] loss: {running_loss / 2000}')
            running_loss = 0.0
print(&quot;Finished training.&quot;)

torch.save(net.state_dict(), './classifier.pth')

For some reason, I'm getting the output
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

at the line: output = net(torch.flatten(X)
Thanks in advance for your help!

When flatten() you remove all dimensions including batch dimension!
Try:
output = net(x.view(x.shape[0], -1))

