Thanks @ptrblck . Please see the attached loss for the above model.

I am also experimenting with one Conv2D layer. For one Conv2D, should I use the following model or what do you suggest?
class FemnistNet(nn.Module):
    def __init__(self):
        super(FemnistNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=5, stride=1, padding=0)
        self.pool1 = nn.MaxPool2d(2, stride=2, )
        self.fc1 = nn.Linear(1024, 512)

    def forward(self, x):
        x = x.view(-1, 1, 28, 28)
        x = self.conv1(x)
        x = th.nn.functional.relu(x)
        x = self.pool1(x)

        x = x.flatten(start_dim=1)
        x = self.fc1(x)
        x = x.softmax()

        return x

Is the above model correct for one Conv2D layer? Moreover, any thoughts on the loss? Gradients/Loss or any other component that needs debugging?
Thanks!