code: def set_device(cuda, local_rank):
      
      if cuda:
          torch.cuda.set_device(local_rank)
          device = torch.device('cuda')
      else:
          device = torch.device('cpu')
      return device
docs: Sets device based on local_rank and returns instance of torch.device.
  
      :param cuda: if True: use cuda
      :param local_rank: local rank of the worker