def METHOD_NAME(cuda, local_rank):
    """
    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
    """
    if cuda:
        torch.cuda.METHOD_NAME(local_rank)
        device = torch.device('cuda')
    else:
        device = torch.device('cpu')
    return device