Error using Custom DistributedSampler
🐛 Bug

When using a custom DistributedSampler, I get the following error while requesting the train dataloader:
File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/accelerators/tpu_backend.py", line 81, in train
    self.tpu_train_in_process(self.trainer.tpu_id, model, self.trainer, self.mp_queue)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/accelerators/tpu_backend.py", line 112, in tpu_train_in_process
    results = trainer.run_pretrain_routine(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 1239, in run_pretrain_routine
    self.train()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 355, in train
    self.reset_train_dataloader(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/data_loading.py", line 205, in reset_train_dataloader
    self.train_dataloader = self.request_dataloader(model.train_dataloader)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/data_loading.py", line 360, in request_dataloader
    dataloader = dataloader_fx()
  File "/content/lib/coreml/coreml/data/data_module.py", line 41, in train_dataloader
    drop_last=False)
  File "/content/lib/coreml/coreml/data/dataloader.py", line 127, in get_dataloader
    'rank': xm.get_ordinal()
  File "/content/lib/coreml/coreml/data/sampler.py", line 139, in __init__
    sampler.dataset, num_replicas, rank, shuffle)
  File "/usr/local/lib/python3.6/dist-packages/torch/utils/data/distributed.py", line 58, in __init__
    rank = dist.get_rank()
  File "/usr/local/lib/python3.6/dist-packages/torch/distributed/distributed_c10d.py", line 598, in get_rank
    _check_default_pg()
  File "/usr/local/lib/python3.6/dist-packages/torch/distributed/distributed_c10d.py", line 210, in _check_default_pg
    "Default process group is not initialized"
AssertionError: Default process group is not initialized
When no custom sampler is used, the code works fine. The DistributedSampler object is created as a wrapper over my actual sampler - the definition can be found here and the object creation happens here.
A quick Google search led me to this PyTorch issue where this comment highlighted that the call to torch.distributed.init_process_group might be missing.
To Reproduce
Use this colab notebook.