import torch


def get_device(gpu=0):
    """
    Enable PyTorch with CUDA if available.
    
    :param int gpu: device number for CUDA
    :returns: device name for CUDA
    :rtype: string
    """
    print("PyTorch version: %s" % torch.__version__)
    dev = torch.device("cuda:{}".format(gpu)) if torch.cuda.is_available() else torch.device("cpu")
    print("Using device: %s" % dev)
    return dev