oneflow.nn.Flatten

class oneflow.nn.Flatten(start_dim: int = 1, end_dim: int = - 1)

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

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

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

For example:

>>> import oneflow as flow
>>> input = flow.Tensor(32, 1, 5, 5)
>>> m = flow.nn.Flatten()
>>> output = m(input)
>>> output.shape
oneflow.Size([32, 25])