Hi,
I have got a task (video analysis) where I have to divide the last fully connected layer of (1+2)D ResNet into multiple chunks and perform classification loss calculation on each individual chunk using CrossEntropyLoss. With more context, I’ve reshaped my model’s last layer (FC) to generate the output of shape (64, 4, 21) where 64 is the batch size, 4 is the number of chunks and 21 is the number of classes to perform classification. So basically my model is computing 4 class values for each video segment. I am computing loss in the following manner:
def calculate_loss(output, target, criterion):
    loss = 0
    for i in range(output.size(1)):
        loss += criterion(output[:, i, :], target[:,i])
    return loss # / output.size(1)

Where I am iterating over each chunk and calculate the classification loss. My model trains well for a few batches (150/2500) but then loss and accuracy both stagnate. And when I check the output of each chunk it produces homogeneous classification values (almost same class), regardless of the target having different classes in each chunk.
Questions:

Is this kind of loss calculation is permissible?
How does the backdrop happen in this scenario?