Hi everyone,
I’m using nn.functional.one_hot() to convert predicted class labels [N×1×H×W] to one-hot encoded tensors [N×C×H×W]. The goal is to calculate loss by a custom loss function.
The problem is nn.functional.one_hot() sets grad_fn = None  so the calculated loss using such encoded predictions fails in backpropagation. How may I retain grad_fn for the one-hot encoded tensors?
Any thoughts are welcomed!