torch.jit.frontend.NotSupportedError: ellipsis is not supported

[lutzroeder]

🐛 Bug

Traceback (most recent call last):
  File "~/repro.py", line 5, in <module>
    torch.jit.script(model).save('alexnet.pt')
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_script.py", line 897, in script
    return torch.jit._recursive.create_script_module(
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 370, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 426, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_script.py", line 388, in _construct
    init_fn(script_module)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 406, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 426, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_script.py", line 388, in _construct
    init_fn(script_module)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 406, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 430, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 322, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 703, in compile_unbound_method
    stub = make_stub(fn, fn.__name__)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 37, in make_stub
    ast = get_jit_def(func, name, self_name="RecursiveScriptModule")
  File "/usr/local/lib/python3.9/site-packages/torch/jit/frontend.py", line 259, in get_jit_def
    return build_def(ctx, fn_def, type_line, def_name, self_name=self_name)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/frontend.py", line 281, in build_def
    param_list = build_param_list(ctx, py_def.args, self_name)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/frontend.py", line 316, in build_param_list
    result = [build_param(ctx, arg, self_name, False) for arg in py_args.args]
  File "/usr/local/lib/python3.9/site-packages/torch/jit/frontend.py", line 316, in <listcomp>
    result = [build_param(ctx, arg, self_name, False) for arg in py_args.args]
  File "/usr/local/lib/python3.9/site-packages/torch/jit/frontend.py", line 326, in build_param
    annotation_expr = build_expr(ctx, py_arg.annotation)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/frontend.py", line 267, in __call__
    return method(ctx, node)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/frontend.py", line 737, in build_Subscript
    raise NotSupportedError(base.range(), "ellipsis is not supported")
torch.jit.frontend.NotSupportedError: ellipsis is not supported:
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/conv.py", line 384
    def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]):
                                                                 ~~~~~~~~ <--- HERE
        if self.padding_mode != 'zeros':
            return F.conv2d(F.pad(input, self._reversed_padding_repeated_twice, mode=self.padding_mode),

To Reproduce

Steps to reproduce the behavior:

pip install --upgrade --pre torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html

import torch
import torchvision
model = torchvision.models.alexnet(pretrained=False)
model.eval()
torch.jit.script(model).save('alexnet.pt')

@malfet @jerryzh168 @raghuramank100

cc @gmagogsfm

[lutzroeder]

#48593

[malfet]

AST tree generation changed a bit in Python-3.9, this fix it: #48676