I’m not sure if this is going to solve your entire problem but one error is that you set the variable i with the 1st dimension of output but then use it to iterate over the third dimension.
you should either do:
for i in range(output.size()[1])
    loss = criterion(output[:, i], labels.type(torch.Longtensor)[i]

or
for i in range(output.size()[3])
    loss = criterion(output[:, :, :, i], labels.type(torch.Longtensor)[:, : i]