Ah ok, thanks for the info.
It looks like a standard segmentation task.
I would suggest to use nn.CrossEntropyLoss for your use case.
Have a look at the following code snippet:
n_class = 10
preds = torch.randn(4, n_class, 24, 24)
labels = torch.empty(4, 24, 24, dtype=torch.long).random_(n_class)

criterion = nn.CrossEntropyLoss()
loss = criterion(preds, labels)

You don’t have to save the target as a “one-hot encoded” tensor, but can just use the class indices for the criterion.
Maybe I misunderstood your question, but what is your goal of implementing the criterion manually?