The first one should work, it would be nice to see the expected output and the one you get.
a = torch.FloatTensor(3,3)
b = a.mean(dim=0) > 2
c = b.nonzero()
a - > tensor([[-2.8721e+27,  4.5780e-41, -2.8721e+27],         [ 4.5780e-41,  0.0000e+00,  0.0000e+00],         [ 0.0000e+00,  6.8929e+34,  8.5771e-39]])
b -> tensor([False,  True, False])
c -> tensor([[1]])