We decided to implement the CNN architecte which looks like this

implement the MNISTNet network architecture
class FashionMNISTNet(nn.Module):
# define the class constructor
def __init__(self):
    
    # call super class constructor
    super(FashionMNISTNet, self).__init__()
    
    # specify convolution layer 1
    self.layer_1 = nn.Sequential(
        nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1),
        nn.BatchNorm2d(32),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2)
    )

    #specify convultion layer 2
    self.layer_2 = nn.Sequential(
        nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3),
        nn.BatchNorm2d(64),
        nn.ReLU(),
        nn.MaxPool2d(2)
    )

    self.linear1 = nn.Linear(64*6*6, 600)

    self.drop = nn.Dropout2d(0.25)

    self.linear2 = nn.Linear(600, 120)
    
    self.linear3 = nn.Linear(120, 10)
    
    # add a softmax to the last layer
    # self.logsoftmax = nn.LogSoftmax(dim=1) # the softmax
    
# define network forward pass
def forward(self, images):
    
    x = self.layer_1(images)

    x = self.layer_2(x)

    x = x.view(x.size(0), -1)

    x = self.linear1(x)

    x = self.drop(x)

    x = self.linear2(x)

    x = self.linear3(x)

    # define layer 3 forward pass
    # x = self.logsoftmax(self.linear3(x))
    
    # return forward pass result
    return x

When we try to run the Network Training with the same code (with the exception of mini_batch_size = 100), all the train_epoch_loss display “nan”:
and the error show this:

image1478×441 28 KB

Does anyone know how to fix this error?
Thank you very much!