Extract element from a matrix

1

Now, I have a float tensor B, that is B = torch.randn(2,3,4), and I also have a bool tensor C, that is the same size as B and its element is true or false. I want to extract the values in B, whose location is the true values of C. What should I do?

2

You could just index B with the BoolTensor:

B = torch.randn(2, 3, 4)
C = torch.randint(0, 2, (2, 3, 4)).bool()
res = B[C]