How to split a Tensor based on class id

[ptrblck]

@AlphaBetaGamma96

Let’s say I had one example that was [0,1,1,0], I’d want a function that returns [0,3] for class_id = 0 or [1,2] for class_id = 1. This is with the idea I could have idx0 and idx1 and then simply do data0 = data[idx0] and data1 = data[idx1]. Is this possible?

Yes, that’s possible with a direct comparison:

x = torch.tensor([0, 1, 1, 0])
idx0 = x == 0
idx1 = x == 1

x[idx0]
# tensor([0, 0])
x[idx1]
# tensor([1, 1])