import torch


def get_device(n_gpu):
    """Returns the name of the PyTorch device to be used, based on the number of available gpus.

    Args:
      n_gpu: number of GPUs available. Use 0 for CPU mode.

    Returns:
      the name of the device to be used by PyTorch.

    """
    return torch.device("cuda:0" if (torch.cuda.is_available() and n_gpu > 0) else "cpu")