def get_device(use_cuda=True, device_id=None, usage=5):
    "Return or set default device; `use_cuda`: None - CUDA if available; True - error if not available; False - CPU"
    if not torch.cuda.is_available():
        use_cuda = False
    else:
        if device_id is None: 
            device_ids = get_available_cuda(usage=usage)
            device_id = device_ids[0]   # get the first available device 
        torch.cuda.set_device(device_id)
    return torch.device(torch.cuda.current_device()) if use_cuda else torch.device('cpu')
Return or set default device; `use_cuda`: None - CUDA if available; True - error if not available; False - CPU
PatchTST_self_supervised/src/basics.py