How to do logical indexing of a torch tensor?

I have a torch tensor x with size torch.Size([10, 1, 28, 28]) and a logical np array a:

array([False, False, False, False, False, False,  True,  True,  True, True])

I want to select the first dimension of x according to logical values in a. But I got error when I try x[a]. What should I do?


1 Like

I got the error can't convert np.ndarray of type numpy.bool_.
But x[torch.from_numpy(np.array(a, dtype=np.uint8))] works. Thank you!

 dim 1 can be like `out = x[:, a]`