I’m sorry I didn’t elaborate on my question better.
To the best of my knowledge, there are two ways of creating one-hot encoded tensors in PyTorch:


.scatter_: with this I can create a one-hot encoded tensor from a given tensor (for me, usually a label map). But for this, I need to define a torch.FloatTensor, and use .scatter_ on it. A torch.FloatTensor needs all its dimensions preemptively defined and does not accept None values. This will usually be an issue when the final mini-batch in a training/validation session is smaller than the other mini-batches. I could alternatively define a new torch.FloatTensor of appropriate shape on each iteration, but I wish to skip this extra operation.

torch.nn.functional.one_hot: with this I can directly yield a one-hot encoded tensor from a given tensor, but the output is in the channels-last format or (N,H,W,C). Hence, torch.transpose is required to convert it into PyTorch’s (N,C,H,W) format. I wish to skip this extra operation as well. (It would be awesome if we have the option in torch.nn.functional.one_hot, to specify which dimension will be the channel dimension.)

Is there a better way? I wish to improve the following rough implementation:
class MyLoss(torch.nn.Module):
    def __init__(self, batch_size, classes):
        super(MyLoss, self).__init__()
        # define some attributes
        self.y_true_one_hot = torch.FloatTensor(batch_size, classes, 240, 240)

    def forward(self, y_pred, y_true):
        with torch.no_grad():
            self.y_true_one_hot.zero_().scatter_(1, y_true, 1)
        # do some operations
        return loss