Check cuda availability before loading models

config.ini

@@ -22,7 +22,7 @@ opt_13b = str:facebook/opt-13b
 [models.params]
 dtype = str:bfloat16
 load_device = str:cpu
-run_device = str:cpu
+run_device = str:cuda
 
 [encrypt.default]
 gen_model = str:gpt2

model_factory.py

@@ -32,6 +32,12 @@ class ModelFactory:
 
             load_device = GlobalConfig.get("models.params", "load_device")
             run_device = GlobalConfig.get("models.params", "run_device")
+            if not torch.cuda.is_available():
+                if load_device == "cuda" or run_device == "cuda":
+                    print("cuda is not available, use cpu instead")
+                    load_device = "cpu"
+                    run_device = "cpu"
+
             if load_device is not None:
                 cls.load_device = torch.device(str(load_device))
             if run_device is not None: