Torch jit script. "Could not export Python function call 'type'" on save method

[vfdev-5]

🐛 Bug

The following error

RuntimeError: 
Could not export Python function call 'type'

is raised only on save and in a specific case. See below.

To Reproduce

import torch


def check_input(img: torch.Tensor) -> bool:
    return isinstance(img, torch.Tensor)


def foo(img: torch.Tensor) -> torch.Tensor:
    if check_input(img):
        return img
    raise TypeError("Unexpected input type. Got {}".format(type(img)))


def bar(img: torch.Tensor) -> torch.Tensor:
    if isinstance(img, torch.Tensor):
        return img
    raise TypeError("Unexpected input type. Got {}".format(type(img)))


if __name__ == "__main__":

    t = torch.rand(3, 12, 12)
    s_bar = torch.jit.script(bar)

    assert bar(t).equal(s_bar(t))
    s_bar.save("test_bar.pt")

    t = torch.rand(3, 12, 12)
    s_foo = torch.jit.script(foo)

    assert foo(t).equal(s_foo(t))
    s_foo.save("test_foo.pt")

and output

Traceback (most recent call last):
  File "test_type.py", line 32, in <module>
    s_foo.save("test_foo.pt")
RuntimeError: 
Could not export Python function call 'type'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__:
  File "test_type.py", line 11
    if check_input(img):
        return img
    raise TypeError("Unexpected input type. Got {}".format(type(img)))
                                                           ~~~~ <--- HERE

Environment

Please copy and paste the output from our
environment collection script
(or fill out the checklist below manually).

You can get the script and run it with:

wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] pytorch-sphinx-theme==0.0.24
[pip3] torch==1.7.0.dev20200928
[pip3] torchvision==0.8.0a0+6f75883
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.1.243             h6bb024c_0  
[conda] mkl                       2020.2                      256  
[conda] mkl-service               2.3.0            py37he904b0f_0  
[conda] mkl_fft                   1.2.0            py37h23d657b_0  
[conda] mkl_random                1.1.1            py37h0573a6f_0  
[conda] numpy                     1.19.2                   pypi_0    pypi
[conda] pytorch                   1.7.0.dev20200928 py3.7_cuda10.1.243_cudnn7.6.3_0    pytorch-nightly
[conda] pytorch-sphinx-theme      0.0.24                    dev_0    <develop>
[conda] torchvision               0.8.0a0+6f75883           dev_0    <develop>

Additional context

Torchvision transforms scripting

cc @gmagogsfm

[gmagogsfm]

This is a duplicate of #42145.

Closing this one, and move all tracking to #42145