torch实现np.nanmean的功能

一条命令搞定

tensor[tensor < float('inf')].mean()

tensor = torch.tensor([1,2,3,4,np.float('nan')])
nanmean = tensor[tensor < float('inf')].mean()
print(tensor)
print(tensor.mean())
print(nanmean)

## output
# tensor([1., 2., 3., 4., nan])
# tensor(nan)
# tensor(2.5000)