[Python] How to convert 2D tensor to 1D tensor

Question

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

Answer 1

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