import torch


def get_device(arguments):
    """
    Sets the device that will be used to training and testing.
    :param arguments: A ArgumentParser Namespace containing "gpu".
    :return: A PyTorch device.
    """

    # Checks if the GPU is available to be used and sets the .
    if arguments.gpu and torch.cuda.is_available():
        return torch.device(f"cuda:{torch.cuda.device_count() - 1}")

    # Sets the device to CPU.
    else:
        return torch.device("cpu")