+2 votes
in Programming Languages by (51.8k points)
How can I flatten a 2D tensor into a 1D tensor?

1 Answer

+2 votes
by (281k points)
selected by
 
Best answer

You can use either flatten() or reshape() to convert a 2D tensor into a 1D tensor.

Using flatten()

>>> import torch
>>> a=torch.tensor([[1,2,3],[4,5,6]])
>>> a
tensor([[1, 2, 3],
        [4, 5, 6]])
>>> a.flatten()
tensor([1, 2, 3, 4, 5, 6])

Using reshape()

>>> a
tensor([[1, 2, 3],
        [4, 5, 6]])

>>> a.reshape(-1,)
tensor([1, 2, 3, 4, 5, 6])


...