Update CodeFormer/basicsr/utils/misc.py

CodeFormer/basicsr/utils/misc.py

@@ -3,6 +3,22 @@ import re
 import random
 import time
 import torch
+def gpu_is_available():
+    return torch.cuda.is_available()
+
+# Note: get_device is also often missing, so it's best practice to add it too.
+def get_device(gpu_id=None):
+    if gpu_id is None:
+        gpu_str = ''
+    elif isinstance(gpu_id, int):
+        gpu_str = f':{gpu_id}'
+    else:
+        raise TypeError('Input should be int value. ')
+    
+    # Prioritize CUDA (GPU) if available, otherwise use CPU
+    if torch.cuda.is_available() and torch.backends.cudnn.is_available():
+        return torch.device(f'cuda{gpu_str}')
+    return torch.device('cpu')
 import numpy as np
 from os import path as osp