Repository: pytorch
Issue: `SyncBatchNorm` doesn't work with subclass of `torch.Tensor`
Labels: oncall: distributed,tensor subclass

### 🐛 Describe the bug

Multi-gpu data parallel pipeline with `SyncBatchNorm` doesn't seem to work with subclasses of `torch.Tensor`. 

To reproduce, save the following script as `test.py` and run it with a basic ddp: `torchrun --nnodes=1 --nproc_per_node=2 test.py` with 2 gpus on 1 node:

```python
import torch
import torch.distributed as dist
from torchvision import models

class SubTensor(torch.Tensor):
    pass

def run():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    x = SubTensor(torch.zeros(1, 3, 128, 128)).to(rank)
    # x = torch.zeros(1, 3, 128, 128).to(rank)  # regular tensor works fine

    mod = models.resnet50(weights=None).to(rank)
    mod = torch.nn.SyncBatchNorm.convert_sync_batchnorm(mod)
    mod = torch.nn.parallel.DistributedDataParallel(mod, device_ids=[rank], output_device=rank)

    out = mod(x)


if __name__ == "__main__":
    run()
```

```
Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1186, in _call_impl
  File "test.py", line 22, in <module>
    run()
  File "test.py", line 18, in run
    out = mod(x)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1186, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1009, in forward
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1009, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 970, in _run_ddp_forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 970, in _run_ddp_forward
    return module_to_run(*inputs[0], **kwargs[0])
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1186, in _call_impl
    return module_to_run(*inputs[0], **kwargs[0])
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1186, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torchvision/models/resnet.py", line 285, in forward
    return self._forward_impl(x)
  File "/opt/conda/lib/python3.8/site-packages/torchvision/models/resnet.py", line 270, in _forward_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torchvision/models/resnet.py", line 285, in forward
    x = self.relu(x)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1186, in _call_impl
    return self._forward_impl(x)
  File "/opt/conda/lib/python3.8/site-packages/torchvision/models/resnet.py", line 270, in _forward_impl
    x = self.relu(x)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1186, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/activation.py", line 102, in forward
    return F.relu(input, inplace=self.inplace)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/functional.py", line 1453, in relu
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/activation.py", line 102, in forward
    return F.relu(input, inplace=self.inplace)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/functional.py", line 1453, in relu
    return handle_torch_function(relu, (input,), input, inplace=inplace)
  File "/opt/conda/lib/python3.8/site-packages/torch/overrides.py", line 1528, in handle_torch_function
    return handle_torch_function(relu, (input,), input, inplace=inplace)
  File "/opt/conda/lib/python3.8/site-packages/torch/overrides.py", line 1528, in handle_torch_function
    result = torch_func_method(public_api, types, args, kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/_tensor.py", line 1089, in __torch_function__
    result = torch_func_method(public_api, types, args, kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/_tensor.py", line 1089, in __torch_function__
    ret = func(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/functional.py", line 1455, in relu
    ret = func(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/functional.py", line 1455, in relu
    result = torch.relu_(input)
RuntimeError: Output 0 of SyncBatchNormBackward is a view and is being modified inplace. This view was created inside a custom Function (or because an input was returned as-is) and the autograd logic to handle view+inplace would override the custom backward associated with the custom Function, leading to incorrect gradients. This behavior is forbidden. You can fix this by cloning the output of the custom Function.
    result = torch.relu_(input)
RuntimeError: Output 0 of SyncBatchNormBackward is a view and is being modified inplace. This view was created inside a custom Function (or because an input was returned as-is) and the autograd logic to handle view+inplace would override the custom backward associated with the custom Function, leading to incorrect gradients. This behavior is forbidden. You can fix this by cloning the output of the custom Function.
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 431) of binary: /opt/conda/bin/python
Traceback (most recent call last):
  File "/opt/conda/bin/torchrun", line 33, in <module>
    sys.exit(load_entry_point('torch', 'console_scripts', 'torchrun')())
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 345, in wrapper
    return f(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/run.py", line 761, in main
    run(args)
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/run.py", line 752, in run
    elastic_launch(
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 132, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 246, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:

```

### Versions

PyTorch version: 1.13.0a0+d321be6
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.23.3
Libc version: glibc-2.31

Python version: 3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:04:10)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-126-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.7.99
GPU models and configuration:
GPU 0: NVIDIA TITAN Xp COLLECTORS EDITION
GPU 1: NVIDIA TITAN Xp COLLECTORS EDITION

Nvidia driver version: 510.85.02
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.5.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.982
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.4
[pip3] pytorch-ignite==0.4.10
[pip3] pytorch-quantization==2.1.2
[pip3] torch==1.13.0a0+d321be6
[pip3] torch-tensorrt==1.2.0a0
[pip3] torchtext==0.11.0a0
[pip3] torchvision==0.14.0a0
[conda] mkl                       2020.4             h726a3e6_304    conda-forge
[conda] mkl-include               2020.4             h726a3e6_304    conda-forge
[conda] numpy                     1.22.4           py38h99721a1_0    conda-forge
[conda] pytorch-ignite            0.4.10                   pypi_0    pypi
[conda] pytorch-quantization      2.1.2                    pypi_0    pypi
[conda] torch                     1.13.0a0+d321be6          pypi_0    pypi
[conda] torch-tensorrt            1.2.0a0                  pypi_0    pypi
[conda] torchtext                 0.11.0a0                 pypi_0    pypi
[conda] torchvision               0.14.0a0                 pypi_0    pypi


see also https://github.com/Project-MONAI/MONAI/issues/5283

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @ezyang