Alan
Alan

Reputation: 15

index a list of tensors

I have a list object named " imgs " of tensors (50 images). I have an array of indices (indi) of length 29.

how do I index the list of tensors with the array of indices?

when I do the following I get:

imgs[indi] 

TypeError: only integer scalar arrays can be converted to a scalar index

Thanks

Upvotes: 0

Views: 634

Answers (1)

jodag
jodag

Reputation: 22184

Assuming these are normal python lists then you can use a list comprehension

result = [imgs[i] for i in indi]

which will give a list of tensors.

If you further want to make this a single tensor containing the images you can use torch.stack

result = torch.stack([imgs[i] for i in indi], dim=0)

Upvotes: 1

Related Questions