oneflow.squeeze

oneflow.squeeze()

This operator removes the specified dimention which size is 1 of the input Tensor. If the dim is not specified, this operator will remove all the dimention which size is 1 of the input Tensor.

The amount of element in return value is the same as Tensor input.

Parameters
  • input (oneflow.Tensor) – the input Tensor.

  • dim (int, optinal) – Defaults to None, if given, the input will be squeezed only in this dimension.

Returns

The result Tensor.

Return type

Tensor

For example:

>>> import oneflow as flow
>>> import numpy as np
>>> input = flow.tensor(np.array([[[[1, 1, 1]]]]).astype(np.int32))
>>> input.shape
oneflow.Size([1, 1, 1, 3])
>>> out = flow.squeeze(input, dim=[1, 2]).shape
>>> out
oneflow.Size([1, 3])