given the following:

dataset = torch.randn((123, 1, 64, 64, 64))
b = torch.mean(dataset, dim=0) > .2
c = b.nonzero()
print(c.shape)


this prints:

torch.Size([3495, 4])

where I expect it to be more like assuming that 23 of the original 123 3d volume elements don’t exceed the threshold

torch.Size(100, 1, 64, 64, 64)