Yeah, I thought so, that’s why I mentioned it again.
I misunderstood the following statement:



 falmasri:

where I want to calculate the softmax of the entire matrix to have a normalized heatmap of the this channel.


You could use the following code:
x = Variable(torch.cat((torch.ones(1, 1, 10, 10), torch.ones(1, 1, 10, 10)*2), dim=0))
softmax = nn.Softmax(dim=1)
y = softmax(x.view(2, -1)).view(2, 1, 10, 10)

Now, each batch sample will have the sum of 1:
print(y[0, 0, ...].sum())
print(y[1, 0, ...].sum())

Does this work for you?