Repository: pytorch
Issue: [JIT] Runtime error when backpropagating through input gradient
Labels: oncall: jit,triaged

## 🐛 Bug

I cannot backpropagate through the grads of one of the inputs of a module that combines two inputs by **concatenating** then.

## To Reproduce

```python
import torch, torch.nn as nn
from torch.autograd import grad

# The following function tests backpropagating through the gradients of one of the inputs of a module
def test_ograd(mod, obs, act):
    obs = obs.clone().requires_grad_(True)
    out = mod(obs, act)
    assert out.grad_fn is not None
    (ograd,) = grad (out.mean(), [obs], create_graph=True)
    assert torch.is_tensor(ograd)
    assert ograd.grad_fn is not None
    ograd.mean().backward()
    assert obs.grad is not None

# This module concats inputs and applies an MLP to then
class Combiner(nn.Module):
    def __init__(self, obs_dim, act_dim):
        super().__init__()
        self.sequential = nn.Sequential(
            nn.Linear(obs_dim + act_dim, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh()
        )

    def forward(self, obs, act):
        output = torch.cat([obs, act], dim=-1)
        output = self.sequential(output)
        return output

# This module ignores the second input and applies an MLP to the first one
class Ignorer(nn.Module):
    def __init__(self, obs_dim):
        super().__init__()
        self.sequential = nn.Sequential(
            nn.Linear(obs_dim, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh()
        )

    def forward(self, obs, act):
        output = self.sequential(obs)
        return output

# This module passes a Linear layer on each input and combines the output before passing the to an MLP
class CombinerV2(nn.Module):
    def __init__(self, obs_dim, act_dim):
        super().__init__()
        self.obs_fc = nn.Linear(obs_dim, 16)
        self.act_fc = nn.Linear(act_dim, 16)
        self.sequential = nn.Sequential(
            nn.Tanh(), nn.Linear(32, 32), nn.Tanh()
        )

    def forward(self, obs, act):
        output = torch.cat([self.obs_fc(obs), self.act_fc(act)], dim=-1)
        output = self.sequential(output)
        return output

class CatAndExp(nn.Module):
    def forward(self, obs, act):
        out = torch.cat([obs, act], dim=-1)
        out = torch.exp(out)
        return out

# Create fake inputs
obs_dim, act_dim = 4, 3
obs, act = torch.randn(1, obs_dim), torch.randn(1, act_dim)

# Test each module
test_ograd(
    Combiner(obs_dim, act_dim),
    obs, act
)
test_ograd(  # This fails
    torch.jit.script(Combiner(obs_dim, act_dim)),
    obs, act
)
test_ograd(
    Ignorer(obs_dim),
    obs, act
)
test_ograd(
    torch.jit.script(Ignorer(obs_dim)),
    obs, act
)
test_ograd(
    CombinerV2(obs_dim, act_dim),
    obs, act
)
test_ograd(  # This fails
    torch.jit.script(CombinerV2(obs_dim, act_dim)),
    obs, act
)
test_ograd(
    CatAndExp(),
    obs, act
)
test_ograd(  # This fails
    torch.jit.script(CatAndExp()),
    obs, act
)
```

## Expected behavior

All tests should pass without fail, however the second, 6th, and 8th ones raise a `RuntimeError` when executing `ograd.mean().backward()`

Example output when running `test_ograd(
    torch.jit.script(Combiner(obs_dim, act_dim)),
    obs, act
)`
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-4-23652e9f96fb> in <module>
----> 1 test_ograd(
      2     torch.jit.script(Combiner(obs_dim, act_dim)),
      3     obs, act
      4 )

<ipython-input-2-ff5ae6d0f669> in test_ograd(mod, obs, act)
     10     assert torch.is_tensor(ograd)
     11     assert ograd.grad_fn is not None
---> 12     ograd.mean().backward()
     13     assert obs.grad is not None
     14

~/Library/Caches/pypoetry/virtualenvs/raylab-5wWf_O6O-py3.8/lib/python3.8/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    183                 products. Defaults to ``False``.
    184         """
--> 185         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    186
    187     def register_hook(self, hook):

~/Library/Caches/pypoetry/virtualenvs/raylab-5wWf_O6O-py3.8/lib/python3.8/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
    123         retain_graph = create_graph
    124
--> 125     Variable._execution_engine.run_backward(
    126         tensors, grad_tensors, retain_graph, create_graph,
    127         allow_unreachable=True)  # allow_unreachable flag

RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "<string>", line 122, in <backward op>
            def backward(grad_outputs: List[Tensor]):
                size = len(grad_outputs)
                grad_self = torch.cat(grad_outputs, dim)
                            ~~~~~~~~~ <--- HERE
                return grad_self, None, None
RuntimeError: dim() called on undefined Tensor
```

It seems that modules that **concatenate** inputs at some point in the forward pass are prone to this error when converted to TorchScript beforehand.


## Environment

PyTorch version: 1.6.0
Is debug build: No
CUDA used to build PyTorch: None

OS: Mac OSX 10.15.6
GCC version: Could not collect
CMake version: version 3.17.3

Python version: 3.8
Is CUDA available: No
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA

Versions of relevant libraries:
[pip3] numpy==1.19.0
[pip3] torch==1.6.0
[conda] Could not collect

## Additional context

<!-- Add any other context about the problem here. -->


cc @suo @gmagogsfm