Thanks for the update!
The previous solution wasn’t working, since your target was one-hot encoded ([ 4, 4]) and you’ve flattened it to a tensor of [16], which created the shape mismatch.
nn.CrossEntropyLoss expects the target to contain the class indices, which can be created via torch.argmax(one_hot_target, dim=1). Your approach with numpy also seems to be correct.