os.environ["CUDA_VISIBLE_DEVICES"] = '0'
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '5678'

torch.distributed.init_process_group(backend='gloo',world_size=1,rank=0)
local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)