vals, idxs = m.max(dim=1, keepdim=True)
print(vals.shape, idxs.shape)   # torch.Size([3, 1]) torch.Size([3, 1])