Solved by theevann in post #3 
              
              
                The torch.sum function will by default remove the dimension on which you apply the sum: 
torch.sum(torch.exp(x), dim=1)  # shape (64)

To keep that dimension, you can use the keepdim optional parameter: 
torch.sum(torch.exp(x), dim=1, keepdim=True)  # shape (64,1)