import torch


def check_gpu(gpu):
    """
    Fuction takes one argument as boolean and provides support for gpu or cpu selection and print out the current device being used.
    Command Line Arguments:
    1. GPU as True value that enables GPU support and use cuda for calculation, and False to enable CPU.
    Function returns device and print out the device being used, either for trianing or inference.
    """
    # If gpu is True gpu is enabled and print out "\nGPU is availabe...". if the gpu didn't exist device switchs to cpu and print out "\nDevice didn't find GPU, using CPU instead"
    if gpu:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        if torch.cuda.is_available():
            print("\nGPU is availabe...")
        else:
             print("\nDevice didn't find GPU, using CPU instead")
    else:
        print("\nCPU is availabe...")
        return torch.device("cpu")

    return device