Flatten

class torch.nn.Flatten(start_dim=1, end_dim=-1)

Flattens a contiguous range of dims into a tensor. For use with Sequential.

Shape:

• Input: $(N, *dims)$

• Output: $(N, \prod *dims)$ (for the default case).

Parameters

• start_dim – first dim to flatten (default = 1).

• end_dim – last dim to flatten (default = -1).

Examples::

>>> input = torch.randn(32, 1, 5, 5)
>>> m = nn.Sequential(
>>>     nn.Conv2d(1, 32, 5, 1, 1),
>>>     nn.Flatten()
>>> )
>>> output = m(input)
>>> output.size()
torch.Size([32, 288])