Use torch.nonzero() as index

3

Thanks! .split makes it easier. As a suggestion if PyTorch tensors could also be indexed by another tensor (of coordinates) that could be even more intuitive than numpy.