Skip to content

Scripting torch.nn.MaxPool fails when return_indices=True #45904

@neginraoof

Description

@neginraoof

🐛 Bug

Type annotation for MaxPool functions does not support the case return_indices=True.
Jit script fails with the exception below:

RuntimeError:
Return value was annotated as having type Tensor but is actually of type Tuple[Tensor, Tensor]:
  File "/opt/anaconda3/lib/python3.8/site-packages/torch/nn/modules/pooling.py", line 227
    def forward(self, input: Tensor) -> Tensor:
        return F.max_pool3d(input, self.kernel_size, self.stride,
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
                            self.padding, self.dilation, self.ceil_mode,
                            ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
                            self.return_indices)
                            ~~~~~~~~~~~~~~~~~~~ <--- HERE

To Reproduce

Steps to reproduce the behavior:

model = torch.nn.MaxPool1d(2, stride=1, return_indices=True)
torch.jit.script(model)

cc @albanD @mruberry @ezyang @malfet @rgommers @xuzhao9 @gramster @gmagogsfm

Metadata

Metadata

Labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions