Can I use heatmap as an output of my model along with the main output in pytorch ? I want to generate a loss from heatmap
You could treat the additional heatmap output in the same way as the original output, i.e. return it in the forward method of your model, and calculate the loss(es). I’m not sure if you are stuck at the implementation of this approach or if you are looking for a general advice regarding loss functions etc.