Repository: pytorch
Issue: JIT fails when calling an exported method inside another one, if kwargs with defaults are not passed
Labels: oncall: jit

## 🐛 Bug

JIT fails on nightly when calling an exported method inside another exported method if kwargs with defaults are not passed

## To Reproduce
```
import torch                                                                                                                                          
from torch import Tensor                                                                                                                              

class Jittable(torch.nn.Module):                                                                                                                      
                                                                                                                                                      
    def forward(self, x: Tensor) -> Tensor:                                                                                                           
        return x                                                                                                                                      
                                                                                                                                                      
    # comment the torch.jit.export to get rid of error:                                                                                               
    @torch.jit.export                                                                                                                                 
    def jit_exportable1(self, x: Tensor, flag: bool = True) -> Tensor:                                                                                
        return x                                                                                                                                      
                                                                                                                                                      
    @torch.jit.export                                                                                                                                 
    def jit_exportable2(self, x: Tensor) -> Tensor:                                                                                                   
        # alternatively, explicitely set y=True to get rid of error                                                                                   
        x = self.jit_exportable1(x)                                                                                                                   
        return x                                                                                                                                      
                                                                                                                                                      
                                                                                                                                                      
model = torch.jit.script(Jittable())                                                                                                                  
                                                                                                                                                      
out1 = model.jit_exportable1(torch.empty(1))                                                                                                          
out2 = model.jit_exportable2(torch.empty(1))
```

This produces the following error:
```
terminate called after throwing an instance of 'torch::jit::ErrorReport'                                                                              
  what():                                                                                                                                             
                                                                                                                                                      
jit_exportable1(__torch__.Jittable self, Tensor x, bool flag) -> (Tensor):                                                                            
Argument flag not provided.                                                                                                                           
:                                                                                                                                                     
  File "test_torch_bug.py", line 18                                                                                                                   
    def jit_exportable2(self, x: Tensor) -> Tensor:                                                                                                   
        # alternatively, explicitely set y=True to get rid of error                                                                                   
        x = self.jit_exportable1(x)                                                                                                                   
            ~~~~~~~~~~~~~~~~~~~~ <--- HERE                                                                                                            
        return x
```

## Expected behavior

Code should not produce an error

## Environment

Collecting environment information...
PyTorch version: 1.9.0.dev20210525+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 8.4.0-3ubuntu2) 8.4.0
Clang version: 10.0.0-4ubuntu1 
CMake version: version 3.16.3

Python version: 3.7 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration: GPU 0: Quadro RTX 6000
Nvidia driver version: 460.73.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.18.1
[pip3] pytorch-memlab==0.2.2
[pip3] torch==1.9.0.dev20210525+cu102
[pip3] torchani==2.3.dev240+g53e5816.d20210525
[pip3] torchaudio==0.9.0.dev20210415
[pip3] torchvision==0.10.0.dev20210415
[pip3] torchviz==0.0.2
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.0.221             h6bb024c_0  
[conda] mkl                       2020.1                      217  
[conda] mkl-service               2.3.0            py37he8ac12f_0  
[conda] mkl_fft                   1.2.0            py37h23d657b_0  
[conda] mkl_random                1.1.1            py37h0573a6f_0  
[conda] numpy                     1.18.1                   pypi_0    pypi
[conda] numpy-base                1.19.2           py37hfa32c7d_0  
[conda] pytorch-memlab            0.2.2                    pypi_0    pypi
[conda] torch                     1.9.0.dev20210525+cu102          pypi_0    pypi
[conda] torchaudio                0.9.0.dev20210415          pypi_0    pypi
[conda] torchvision               0.10.0.dev20210415          pypi_0    pypi
[conda] torchviz                  0.0.2                    pypi_0    pypi


## Additional context
For reference, this came up because in my code I have multiple entry points for my model, and some of them call each other, so all the entry points are decorated with `torch.jit.export`.

This issue can be easily bypassed by not exporting the method `jit_exportable1`. The function is correctly scripted anyways since it is called from inside another exported function, however, the unclear error message together with the fact that actually specifying the optional kwargs makes this code work makes me suspect this kind of code is supposed to be supported and this is a bug.

Edit: changed `collect_env.py` output since I tested with the latest nightly just in case but the error is still there

cc @gmagogsfm