Indexing 4D Tensor With 2D indices

Suppose I have a tensor, a, that I wish to index.

a has:
B batches;
C channels;
and dimensions X, Y, Z that contain the values I want to index.

>> torch.Size([B, C, X, Y, Z])

I have a list of tensors, b, of length B, each containing N indices I want to obtain. Each element of the list may have different values for N ie. N0, N1

>> B

>> torch.Size([N0, 3])

>> torch.Size([N1, 3])

To be clear, having indexed the tensor, the final result I want is the following:

>> B

>> torch.Size([C, N0])

>> torch.Size([C, N1])

What is the best way of going about this?

I’m unsure how to understand the “values” of N0 and N1. Could you post an example code snippet with real shapes and values (even if it’s a slow nested loop)?