@KFrank I had a look over the “K-dimensional case” in torch.nn.CrossEntropyLoss 30 and it really helped solve my problem. What I was essentially doing can be done with criterion = torch.nn.CrossEntropyLoss(reduction='sum').
Thanks for the help!