Remove certain elements of all the tensors in a list of dictionary of tensors in pytorch
I am new to pytorch, and trying to learn by playing with simple codes. What I have is a list of length zero, populated by a dictionary, where values in this dictionary are tensors. This is what this list looks like: 

A = [{'boxes': tensor([[ 142.1232,  142.9373, 1106.0452,  971.3792],
    [ 259.1277,  618.4834, 1100.1293, 1028.8989],
    [ 232.1346,  692.5888,  763.3408, 1028.6766],
    [ 206.8070,  312.2080, 1137.1434, 1013.4373],
    [ 495.9471,  675.7287,  978.5932, 1012.7568]], grad_fn=&lt;StackBackward&gt;), 'labels': tensor([16,  1,  1,  1,  1]), 'scores': tensor([0.9988, 0.9489, 0.5228, 0.3500, 0.0639], grad_fn=&lt;IndexBackward&gt;), 'masks': tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]]], grad_fn=&lt;UnsqueezeBackward0&gt;)}]


This list is the output of Mask-RCNN and I want to remove certain element of all the tensors in the nested dictionary. In this case, I only want to keep information related with class label "1". Information related to each class ('boxes', 'labels', 'scores', and 'masks') are all in the same location(s) (index) in each of the tensors. So, I find the index of all the "1"s in the tensor with key "labels":

idxOfClass = [i for i, x in enumerate(list(pred[0]['labels'])) if x == 1]


which gives me: [1, 2, 3, 4]. 
Then, I want to keep all the values located at indices in idxOfClass in all of the tensors in the nested dictionary. If I do something like this:

Anew = [{pred[0]['boxes'][idxOfClass],pred[0]['labels'][idxOfClass],pred[0]['masks'][idxOfClass],pred[0]['scores'][idxOfClass]}]


I get:

[{tensor([[ 259.1277,  618.4834, 1100.1293, 1028.8989],
    [ 232.1346,  692.5888,  763.3408, 1028.6766],
    [ 206.8070,  312.2080, 1137.1434, 1013.4373],
    [ 495.9471,  675.7287,  978.5932, 1012.7568]], grad_fn= 
    &lt;IndexBackward&gt;), tensor([0.9489, 0.5228, 0.3500, 0.0639], 
    grad_fn= 
    &lt;IndexBackward&gt;), tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]]], grad_fn=&lt;IndexBackward&gt;), 
    tensor([1, 1, 1, 1])}]


But, this is not a list of dictionary with tensors as the values of the dictionary. This is a list of tensors, with no key-value structure of the nested dictionary. My question is, "is there any way to keep the original structure of the list when I remove certain elements of all the tensors using indices of specific elements?".

You can just add key names when you construct the new pred res. 

Anew = [{'boxes': pred[0]['boxes'][idxOfClass],'labels': pred[0]['labels'][idxOfClass],'masks': pred[0]['masks'][idxOfClass],'scores': pred[0]['scores'][idxOfClass]}]

