import torch


def activate_gpu(gpu='GPU'):
    """Use GPU if available and requested by user. Defaults to use GPU if available."""
    if torch.cuda.is_available() and gpu.lower() == 'gpu':
        print('Running on GPU')
        device = torch.device('cuda:0')
    else:
        print('Running on CPU')
        device = torch.device('cpu')
        
    return device