How to get columns from 2D tensor list in Pytorch

Question

I have a one 2D list that contains tensors like:

[
 [tensor([-0.0705,  1.2019]), tensor([[0]]), tensor([-1.3865], dtype=torch.float64), tensor([-0.0744,  1.1880]), tensor([False])],
 [tensor([-0.0187,  1.3574]), tensor([[2]]), tensor([0.3373], dtype=torch.float64), tensor([-0.0221,  1.3473]), tensor([False])],
 [....] ]

The outer list contains 64 little list. One little list contains 5 different tensor elements.

And I want to get first elements of inner lists like tensor([-0.0705, 1.2019]) and tensor([-0.0187, 1.3574]) and create tensor like 64x2 to feed my neural net.

How can I do this in the fastest way?

Thanks

Answer 1

How about using slices?

import torch
import numpy as np
x = [
 [torch.tensor([-0.0705,  1.2019]), torch.tensor([0]), torch.tensor([-1.3865], dtype=torch.float64), torch.tensor([-0.0744,  1.1880]), torch.tensor([False])],
 [torch.tensor([-0.0187,  1.3574]), torch.tensor([2]), torch.tensor([0.3373], dtype=torch.float64), torch.tensor([-0.0221,  1.3473]), torch.tensor([False])]]
x = list(map(lambda x:list(map(lambda z:z.tolist(), x)), x))
print(x)
x = np.array(x)[:, 0]
x = list(map(lambda z:torch.tensor(z), x))
print(x)

Answer 2

Use list comprehension

[item[0] for item in your_list]

Example:

li = [[tensor([-0.0705,  1.2019]), tensor([[0]]), tensor([-1.3865], dtype=torch.float64), tensor([-0.0744,  1.1880]), tensor([False])],
     [tensor([-0.0187,  1.3574]), tensor([[2]]), tensor([0.3373], dtype=torch.float64), tensor([-0.0221,  1.3473]), tensor([False])]]

[item[0] for item in li]
[tensor([-0.0705,  1.2019]), tensor([-0.0187,  1.3574])]