How to extract tensors to numpy arrays or lists from a larger pytorch tensor

Question

I have a list of pytorch tensors as shown below:

data = [[tensor([0, 0, 0]), tensor([1, 2, 3])],
        [tensor([0, 0, 0]), tensor([4, 5, 6])]]

Now this is just a sample data, the actual one is quite large but the structure is similar.

Question: I want to extract the tensor([1, 2, 3]), tensor([4, 5, 6]) i.e., the index 1 tensors from data to either a numpy array or a list in flattened form.

Expected Output:

out = array([1, 2, 3, 4, 5, 6])

OR

out = [1, 2, 3, 4, 5, 6]

---

I have tried several ways one including map function like:

map(lambda x: x[1].numpy(), data)

This gives:

[array([1, 2, 3]),
 array([4, 5, 6])]

And I'm unable to get the desired result with any other method I'm using.

Answer 1

OK, you can just do this.

out = np.concatenate(list(map(lambda x: x[1].numpy(), data)))

Answer 2

You can convert a nested list of tensors to a tensor/numpy array with a nested stack:

data = np.stack([np.stack([d for d in d_]) for d_ in data])

You can then easily index this, and concatenate the output:

>>> np.concatenate(data[:,1])
array([[1, 2, 3],
       [4, 5, 6]])