m = torch.tensor([[1., 2., 3.],
                  [4., 5., 6.],
                  [7., 8., 9.]])          # (3, 3)

wrong   = m / m.sum(dim=1)                 # (3, 3) / (3,)  — no error!
correct = m / m.sum(dim=1, keepdim=True)   # (3, 3) / (3, 1)

print(wrong.sum(dim=1))    # tensor([0.4250, 1.2500, 2.0750])  ← NOT 1s ❌
print(correct.sum(dim=1))  # tensor([1., 1., 1.])              ← correct ✅