Repository: pytorch
Issue: jit tracer doesn't work with unflatten layer
Labels: oncall: jit,weeks

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:

1. Run the following code:

```
import torch                                                                                                                                                                                                       
import torch.nn                                                                                                                                                                                                    
import torch.onnx                                                                                                                                                                                                  
                                                                                                                                                                                                                   
model = torch.nn.Sequential(                                                                                                                                                                                       
    torch.nn.Unflatten(1, (1, 10)),                                                                                                                                                                                
)                                                                                                                                                                                                                  
                                                                                                                                                                                                                   
# Create the right input shape (e.g. for an image)                                                                                                                                                                 
dummy_input = torch.randn(100, 10)                                                                                                                                                                                 
                                                                                                                                                                                                                   
print(torch.__version__)                                                                                                                                                                                           
torch.onnx.export(model, dummy_input, "onnx_model_name.onnx", verbose=True)    
```

This outputs: 

```
1.7.1
Traceback (most recent call last):
  File "test_export.py", line 13, in <module>
    torch.onnx.export(model, dummy_input, "onnx_model_name.onnx", verbose=True)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/onnx/__init__.py", line 230, in export
    custom_opsets, enable_onnx_checker, use_external_data_format)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/onnx/utils.py", line 91, in export
    use_external_data_format=use_external_data_format)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/onnx/utils.py", line 639, in _export
    dynamic_axes=dynamic_axes)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/onnx/utils.py", line 411, in _model_to_graph
    use_new_jit_passes)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/onnx/utils.py", line 379, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/onnx/utils.py", line 342, in _trace_and_get_graph_from_model
    torch.jit._get_trace_graph(model, args, strict=False, _force_outplace=False, _return_inputs_states=True)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/jit/_trace.py", line 1148, in _get_trace_graph
    outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/jit/_trace.py", line 130, in forward
    self._force_outplace,
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/jit/_trace.py", line 116, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 725, in _call_impl
    result = self._slow_forward(*input, **kwargs)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 709, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/nn/modules/container.py", line 117, in forward
    input = module(input)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 725, in _call_impl
    result = self._slow_forward(*input, **kwargs)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 709, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/nn/modules/flatten.py", line 133, in forward
    return input.unflatten(self.dim, self.unflattened_size)
  File "/home/justin/anaconda3/lib/python3.7/site-packages/torch/tensor.py", line 861, in unflatten
    return super(Tensor, self).unflatten(dim, sizes, names)
RuntimeError: NYI: Named tensors are not supported with the tracer
```

## Expected behavior

I'd expect either no error at all, or one specifically saying why unflatten isn't supported -- there are no named tensors here as far as I can tell. 

## Environment

```
Collecting environment information...
PyTorch version: 1.7.1
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Pop!_OS 20.04 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.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: Could not collect
GPU models and configuration: GPU 0: GeForce RTX 2070
Nvidia driver version: 455.38
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.17.4
[pip3] numpydoc==1.1.0
[pip3] torch==1.7.1
[pip3] torchsummary==1.5.1
[pip3] torchvision==0.8.1
[pip3] torchviz==0.0.1
[conda] blas                      1.0                         mkl    anaconda
[conda] cudatoolkit               11.0.221             h6bb024c_0    anaconda
[conda] mkl                       2019.4                      243    anaconda
[conda] mkl-service               2.3.0            py37he904b0f_0    anaconda
[conda] mkl_fft                   1.2.0            py37h23d657b_0    anaconda
[conda] mkl_random                1.1.0            py37hd6b4f25_0  
[conda] numpy                     1.19.1           py37hbc911f0_0    anaconda
[conda] numpy-base                1.19.1           py37hfa32c7d_0    anaconda
[conda] numpydoc                  1.1.0                      py_0    anaconda
[conda] torch                     1.7.0                    pypi_0    pypi
[conda] torchsummary              1.5.1                    pypi_0    pypi
[conda] torchvision               0.8.1                    pypi_0    pypi
[conda] torchviz                  0.0.1                    pypi_0    pypi
```

## Additional context

A custom module with reshape in it works fine as a workaround it seems like. 

cc @gmagogsfm