Data in GPU selected by a dictionary (in CPU)

Hello,

The following example seems very simple but when I looked at it I cannot figure out what the data workflow would look like. Since I am new to Pytorch and I am careful about efficiency, I like to understand it in a good way so I feel confident to proceed. Thank you in advance.

For example, if you have a tensor already in GPU

X = torch.from_numpy(np.array([1,2,3,4,5]))
Xg = X.to(device)

After this I was confirmed that Xg is indeed in GPU as device='cuda:0'

Now I have a use case: I have a Pythonic dictionary like

index_map = {"a": np.array([0,1]), "b":np.array([2,3,4])}

and I use this dictionary to select corresponding values inside Xg, like

Xg[index_map["a"]]

and I get the result

tensor([1, 2], device='cuda:0')

It looks obvious, but the question of me is that since index_map is a python dictionary, and from a few other topics and experiment, a dictionary cannot be directly moved to GPU as a whole. So I am wondering when this last operation is taken, since Xg is in GPU but index_map is not, how does this work? Does the workflow like (1) evaluate index_map[“a”] in CPU and get the corresponding index (2) pass those indices to GPU and operate on Xg ? Or, index_map[“a”] is somehow silently copied into the GPU beforehand, and all stuffs are inside GPU?

The indices should be copied to the device before the actual indexing kernel is called, if I’m not mistaken.