Sorry, here’s more code, I’m following guides from tutorials and the documentation.
def main(rank, args):
  dist.init_process_group(...)
  # ... load data ...
  train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset.lengths(), 
    num_replicas = args.gpus, 
    rank = rank, 
    shuffle = True
  )
  train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size = args.batch_size,
    num_workers = args.data_workers,
    collate_fn = mybatchifyfunc,
    pin_memory = True,
    drop_last = True,
    shuffle = False
  )

  # ... machine learning ...

if __name__ == "__main__":
  # ... set up logging, parse args...
  torch.multiprocessing.spawn(main, args=(args,), nprocs=args.gpus, join=True)