pytorch
543b7fb9 - [JIT] Fix type annotations of pooling modules (#65847)

Commit
3 years ago
[JIT] Fix type annotations of pooling modules (#65847) Summary: All of the pooling modules except MaxUnpool and LPPool return either a Tensor or [Tensor, Tensor]. The current type annotations are inaccurate, and prevent scripting the module if return_indices is set as True in the module. There's not a great way to make this agree with mypy because the overload is dependent on the value of return_indices, an attribute. I tried changing the annotations from `Tensor` to `Union[Tensor, Tuple[Tensor, Tensor]]`, but that breaks a bunch of uses that have return_indices=False. For example, this breaks: https://github.com/pytorch/pytorch/blob/4e94e84f651f7fefddd34212423a129bf56d1002/torch/nn/modules/container.py#L139 Also clean up how test names were being constructed in test_jit, since otherwise we were getting name collisions when there were two tests on the same nn.Module. Fixes https://github.com/pytorch/pytorch/issues/45904 Pull Request resolved: https://github.com/pytorch/pytorch/pull/65847 Reviewed By: ZolotukhinM Differential Revision: D31462517 Pulled By: eellison fbshipit-source-id: 6f9e8df1be6c75e5e1e9bae07cf3ad3603ba59bd
Author
Parents
Loading