How to select subtensor from a tensor
I have a tensor A with the dimension of [N1, N2, N3/2, 2, N4, N5]. Another tensor B is an boolean index tensor with the dimension of [N1, N2, N3/2] = 1,0,0,0,1....
Now, I want to get a subtensor C with the dimension of [N1, N2, N3/2, N4, N5] using the index of B .
How could I implement this with pytorch?
Note that I don't incline to use for-loop because it is very slow.

I have looked up some functions in pytorch but found none of them are suitable for my task

You are making no sense: How can C have the same shape as A when B only selects part of the entries of A to copy to C??

Suppose B has k non-zero elements (that is, B selects k elements out of N1*N2*N3/2 matrices of size N4*N5 in A) then C can have shape of [k, N4, N5] only with k &lt; N1*N2*N3/2. 

You can use nonzero() to convert B from logical indices to k integer indices and then use these to select the k elements

C = A.view(-1, *A.shape[-2:])[B.nonzero(), ...]

