I was trying to reproduce this test for an input with size torch.Size([100, 9, 9, 4]) and target with size torch.Size([100, 9, 9]), but got error.

RuntimeError: Assertion `input0 == target0 && input2 == target1 && input3 == target2’ failed.

I was wondering what is the best way to support batch size, etc. for computing the loss. In my attempt, I used input.view(-1,4) with target.view(-1,4).argmax(). Is there a better way to do it?