How can I flatten a 2D tensor into a 1D tensor?

+2 votes

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