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

1 Answer

+1 vote
by (48.9k points)

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])


...