torch.Tensor.flatten

Tensor.flatten(start_dim=0, end_dim=-1) → Tensor

See torch.flatten()