I have a multi dimensional output model with the shape of (B,C,T) before the softmax layer. Its target is a row wise one hot encoded matrix with the same shape  of model prediction ie (B,C,T) . The trouble is PyTorch softmax method doesn’t working for row wise one hot encoded values. I wrote this sample code to show that the output value after the softmax layer is not row summed to one.
import torch

# shape (Batch, Class, Time Step)
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.nn.Softmax(1)(target)

assert torch.all(target[0].sum(1).to(torch.float)==input[0].sum(1).to(torch.float)),"Failed: sum is not equal to one"

loss = torch.nn.BCELoss()(input,target)

What is a suitable activation and loss function for this kind of multi dimensional one hot encoded data?