Solved by albanD in post #8 
              
              
                Ho ok 
So you want out.sum(-1).sum(-1) to be 1! 
Then you need to collapse the last dims together before the sofmax: 
# Variables don't exist anymore, you can remove
# them and just use Tensors everywhere 
A = torch.rand(1, 2, 3, 3)

A_view = A.view(1, 2, -1)
out_view = F.softmax(A, dim=-1)
out = ou…