nn.CrossEntropyLoss expects raw logits, so you should not use a softmax on the input tensor.
Since the one-hot encoding is in dim2 in the target tensor and deals with 4 classes, your input should thus have the shape [3, 4, 3].
This code snippet shows its usage:
target = torch.tensor([[[0,1,0,0],
                        [1,0,0,0],
                        [0,0,0,1]],
                        
                        [[1,0,0,0],
                        [0,0,1,0],
                        [0,1,0,0]],
                         
                        [[0,0,1,0],
                        [0,0,0,1],
                        [1,0,0,0]]],dtype=torch.float)

input = torch.randn(3, 4, 3)
criterion = nn.CrossEntropyLoss()
target = torch.argmax(target, dim=2)
loss = criterion(input, target)