Could you try to add following codes in your script?
torch.distributed.init_process_group(backend=‘YOUR BACKEND’,
init_method=‘env://’)
torch.cuda.set_device(args.local_rank)  # before your code runs
also python -m torch.distributed.launch --help to understand more about how to use torch.distributed.launch