Repository: pytorch
Issue: torchscript does not work with `SyncBatchNorm` layers
Labels: oncall: jit,Stale

### 🐛 Describe the bug

When attempting to script a model that has sync-batchnorm configured, an error is produced. 

## To Reproduce

```py
import torch
import torch.distributed
import os


class BoringModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)
        self.batch_norm = torch.nn.BatchNorm1d(2)

    def forward(self, x):
        return self.batch_norm(self.layer(x))


def run(local_rank):
    device = torch.device("cuda", local_rank)
    model = BoringModel()

    torch.distributed.init_process_group("gloo", world_size=2, rank=local_rank)
    torch.cuda.set_device(local_rank)

    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
    model = model.to(device)
    script = torch.jit.script(model)


if __name__ == "__main__":
    local_rank = int(os.environ["LOCAL_RANK"])
    run(local_rank)

```
Run the script with this command:
```
torchrun --standalone --nnodes 1 --nproc_per_node=2 pl_examples/bug_report/bug.py
```

Traceback:
```

Traceback (most recent call last):
  File "/home/adrian/repositories/pytorch-lightning/pl_examples/bug_report/pure_python.py", line 30, in <module>
    run(local_rank)
  File "/home/adrian/repositories/pytorch-lightning/pl_examples/bug_report/pure_python.py", line 25, in run
    script = torch.jit.script(model)
  File "/home/adrian/repositories/pytorch/torch/jit/_script.py", line 1258, in script
    return torch.jit._recursive.create_script_module(
  File "/home/adrian/repositories/pytorch/torch/jit/_recursive.py", line 451, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/home/adrian/repositories/pytorch/torch/jit/_recursive.py", line 513, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/home/adrian/repositories/pytorch/torch/jit/_script.py", line 587, in _construct
    init_fn(script_module)
  File "/home/adrian/repositories/pytorch/torch/jit/_recursive.py", line 491, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/home/adrian/repositories/pytorch/torch/jit/_recursive.py", line 517, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/home/adrian/repositories/pytorch/torch/jit/_recursive.py", line 368, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError:
Could not cast value of type NoneType to bool:
  File "/home/adrian/repositories/pytorch/torch/nn/modules/batchnorm.py", line 730
        if need_sync:
            process_group = torch.distributed.group.WORLD
            if self.process_group:
               ~~~~~~~~~~~~~~~~~~ <--- HERE
                process_group = self.process_group
            world_size = torch.distributed.get_world_size(process_group)

ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 171033) of binary: /home/adrian/anaconda3/envs/pytorch/bin/python
Traceback (most recent call last):
  File "/home/adrian/anaconda3/envs/pytorch/bin/torchrun", line 33, in <module>
    sys.exit(load_entry_point('torch', 'console_scripts', 'torchrun')())
  File "/home/adrian/repositories/pytorch/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 345, in wrapper
    return f(*args, **kwargs)
  File "/home/adrian/repositories/pytorch/torch/distributed/run.py", line 724, in main
    run(args)
  File "/home/adrian/repositories/pytorch/torch/distributed/run.py", line 715, in run
    elastic_launch(
  File "/home/adrian/repositories/pytorch/torch/distributed/launcher/api.py", line 131, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/home/adrian/repositories/pytorch/torch/distributed/launcher/api.py", line 245, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
============================================================
pl_examples/bug_report/pure_python.py FAILED
------------------------------------------------------------
Failures:
[1]:
  time      : 2021-12-15_08:08:28
  host      : lambda-server4
  rank      : 1 (local_rank: 1)
  exitcode  : 1 (pid: 171034)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2021-12-15_08:08:28
  host      : lambda-server4
  rank      : 0 (local_rank: 0)
  exitcode  : 1 (pid: 171033)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
```

## Expected Behavior

Option 1: torchscript ignores the wrapped module
Option 2: The user is expected to unwrap the batchnorm layers. Feature request here: #41081


## Additional Context

- This issue was reported in PyTorch Lightning: https://github.com/PyTorchLightning/pytorch-lightning/issues/11076
- A related feature request:  #41353


### Versions

```

Collecting environment information...
PyTorch version: 1.11.0a0+gita31aea8
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.9.7 (default, Sep 16 2021, 13:09:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-58-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.1.105
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090
GPU 2: NVIDIA GeForce RTX 3090
GPU 3: NVIDIA GeForce RTX 3090
GPU 4: NVIDIA GeForce RTX 3090
GPU 5: NVIDIA GeForce RTX 3090
GPU 6: NVIDIA GeForce RTX 3090
GPU 7: NVIDIA GeForce RTX 3090

Nvidia driver version: 470.57.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.4
[pip3] pytorch-lightning==1.6.0.dev0
[pip3] torch==1.11.0a0+gita31aea8
[pip3] torchmetrics==0.6.1
[conda] numpy                     1.21.4                   pypi_0    pypi
[conda] pytorch-lightning         1.6.0.dev0                dev_0    <develop>
[conda] torch                     1.11.0a0+gita31aea8           dev_0    <develop>
[conda] torchmetrics              0.6.1                    pypi_0    pypi
```