Why does softmax cause my nn to not converge

Question

I have a pytorch model made up of a several convolutional and groupnorm layers which eventually feed into fully connected and eventually a softmax. With the softmax, the model never converges and plateaus quickly (tried up to epoch > 300). But removing the softmax causes the model to decrease in loss quickly and converge within 50 epochs.

My understanding is that multi-class networks should receive a softmax layer followed by categorical cross entropy loss, why in this case does the model only converge after a fully connected layer?

The inputs to the model are encoded videos of size (128,64,64) all normalized between 0 and 1.

class ActionClass(nn.Module):
def __init__(self, class_num):
    super(ActionClass, self).__init__()

    self.classNum = class_num

    self.conv_b1_1 = nn.Conv2d(161, 128, kernel_size=3, stride=2, padding=1, groups=1, bias=False, dilation=1)
    self.groupNorm_b1_1 = nn.GroupNorm(128, 128)
    self.conv_b1_2 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1)
    self.groupNorm_b1_2 = nn.GroupNorm(256, 256)

    self.conv_b2_1 = nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1, groups=1, bias=False, dilation=1)
    self.groupNorm_b2_1 = nn.GroupNorm(256, 256)
    self.conv_b2_2 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1)
    self.groupNorm_b2_2 = nn.GroupNorm(512, 512)

    self.conv_b3_1 = nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1, groups=1, bias=False, dilation=1)
    self.groupNorm_b3_1 = nn.GroupNorm(512, 512)
    self.conv_b3_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1)
    self.groupNorm_b3_2 = nn.GroupNorm(512, 512)

    self.relu = nn.LeakyReLU(inplace=True)
    self.softmax = nn.Softmax(dim=0)
    self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
    self.fc = nn.Linear(512, self.classNum)

def forward(self, x):
    x = self.conv_b1_1(x)
    x = self.groupNorm_b1_1(x)

    x = self.relu(x)

    x = self.conv_b1_2(x)
    x = self.groupNorm_b1_2(x)
    x = self.relu(x)

    x = self.conv_b2_1(x)
    x = self.groupNorm_b2_1(x)
    x = self.relu(x)

    x = self.conv_b2_2(x)
    x = self.groupNorm_b2_2(x)
    x = self.relu(x)

    x = self.conv_b3_1(x)
    x = self.groupNorm_b3_1(x)
    x = self.relu(x)

    x = self.conv_b3_2(x)
    x = self.groupNorm_b3_2(x)
    x = self.relu(x)

    x = self.avgpool(x)
    x = x.flatten(start_dim=1)
    x = self.fc(x)
    x = self.softmax(x)

    return x

Answer 1

The reason that this didn't work is Pytorch's implementation of cross entropy loss in nn.CrossEntropyLoss expects logits, not the probabilities output by softmax as suggested in shimao's comment.