Repository: pytorch
Issue: torch.jit.trace() fails on a GCN with sparse inputs and dense layers
Labels: oncall: jit

## 🐛 Bug

I'm trying to trace a basic GCN with sparse input tensors and dense intermediate layers with jit. I have included the minimum code required to replicate these failures under **To Reproduce**.

Tracing gives the following warnings:
```
[W pybind_utils.h:521] Warning: Using sparse tensors in TorchScript is experimental. Many optimization pathways have not been thoroughly tested with sparse tensors. Please include the fact that the network is running sparse tensors in any bug reports submitted. (function operator ())
[W pybind_utils.h:851] Warning: Using sparse tensors in TorchScript is experimental. Many optimization pathways have not been thoroughly tested with sparse tensors. Please include the fact that the network is running sparse tensors in any bug reports submitted. (function operator ())
```
before failing with the following exception:
```
Traceback (most recent call last):
  File ".\minimum_example.py", line 55, in <module>
    traced = torch.jit.trace(model, inputs)
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\jit\_trace.py", line 733, in trace
    return trace_module(
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\jit\_trace.py", line 958, in trace_module
    _check_trace(
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\autograd\grad_mode.py", line 26, in decorate_context
    return func(*args, **kwargs)
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\jit\_trace.py", line 327, in _check_trace
    copied_dict[name] = _clone_inputs(data)
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\jit\_trace.py", line 158, in _clone_inputs
    return function._nested_map(
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\autograd\function.py", line 282, in _map
    return type(obj)(mapped)
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\autograd\function.py", line 278, in <genexpr>
    mapped = (_map(x) for x in obj)
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\autograd\function.py", line 274, in _map
    return fn(obj)
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\jit\_trace.py", line 148, in clone_input
    a.detach()
RuntimeError: unsupported memory format option Preserve
```

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:

1. Construct a basic gcn (example code below)
```
import torch

### LAYER and GCN DEFINITION

class HiddenLayerSparse(torch.nn.Module):
    def __init__(self, dimf_in, dimf_out):
        super().__init__()
        self.weights = torch.nn.Parameter(torch.rand(dimf_in, dimf_out, dtype=float))

    def forward(self, adj, x):
        x = torch.mm(x, self.weights)
        x = torch.mm(adj, x)
        return x

class ExampleSparseGCN(torch.nn.Module):
    def __init__(self, sizes):
        super().__init__()
        self.sizes = sizes
        self.hidden_layers = torch.nn.ModuleList([HiddenLayerSparse(sizes[i], sizes[i+1]) for i in range(len(sizes)-1)])
        self.nonlinear = torch.nn.ReLU()
        self.softmax = torch.nn.Softmax(dim=1)

    def forward(self, adj, x):
        for h in self.hidden_layers:
            x = self.nonlinear(h(adj, x))
        return self.softmax(x)

### EXAMPLE SPARSE INPUTS

adjacency = torch.tensor([[1,0],[0,1]], dtype=float).to_sparse()
features = torch.tensor([[1,0,0,1,0],[0,1,0,0,1]],dtype=float).to_sparse()
labels = torch.tensor([[1,0,0], [0,1,0]],dtype=float)
inputs = (adjacency, features)

model = ExampleSparseGCN([5, 4, 3]) # training does not fail with [5,3]
loss_fn =  torch.nn.L1Loss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

_train = False      # FLAG TO ENABLE TRAINING
_trace = True      # FLAG TO ENABLE TRACING

# TRAINING 
if _train:
    model.train()
    for t in range(50):
        optimizer.zero_grad()
        output = model(*inputs)
        loss = loss_fn(output, labels)
        loss.backward()
        optimizer.step() 
    print('Train Complete')

# TRACING
if _trace:
    model.eval()
    with torch.no_grad():
        #torch.onnx.export(model, inputs, path, do_constant_folding=True, verbose=True, export_params=True)
        traced = torch.jit.trace(model, inputs)
    print('Trace Complete')
```
2. Attempt to trace a model without training. (`python example.py`)
3. Attempt to trace a model with training (set `_train = True`)

Any GCN with sparse inputs **or** dense inputs with `torch.sparse.mm` fails to trace. Using `torch.sparse.mm` breaks training with exception:
```
Traceback (most recent call last):
  File ".\minimum_example.py", line 46, in <module>
    loss.backward()
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\tensor.py", line 221, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\autograd\__init__.py", line 130, in backward
    Variable._execution_engine.run_backward(
RuntimeError: sparse_.is_sparse() INTERNAL ASSERT FAILED at "..\\torch\\csrc\\autograd\\FunctionsManual.cpp":560, please report a bug to PyTorch.
```

## Expected behavior

I expect the model to train and trace without triggering any exceptions using `torch.sparse.mm`. Currently tracing always fails, and using `torch.sparse.mm` causes training to fail in the backward pass (which can be alleviated by using `torch.mm`.)

## Environment

```
PyTorch version: 1.7.1+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Education
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect

Python version: 3.8 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: GeForce RTX 2070
Nvidia driver version: 456.71
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.18.5
[pip3] numpydoc==1.1.0
[pip3] torch==1.7.1
[conda] blas                      1.0                         mkl
[conda] mkl                       2020.1                      216
[conda] mkl-service               2.3.0            py38hb782905_0
[conda] mkl_fft                   1.1.0            py38h45dec08_0
[conda] mkl_random                1.1.1            py38h47e9c7a_0
[conda] numpy                     1.18.5           py38h6530119_0
[conda] numpy-base                1.18.5           py38hc3f5095_0
[conda] numpydoc                  1.1.0                      py_0
[conda] torch                     1.7.1                    pypi_0    pypi
```

## Additional context

Please let me know if more context/information/support is needed from my end.


cc @gmagogsfm