Merge pull request #66 from Ivan-Inby/bugfix/fallback-to-cpu-if-no-cuda

fix: fallback to CPU when CUDA is not available
2025-06-15 22:05:53 +08:00
parent 8ac1495de8 3b86dc6254
commit cb1c409cea
1 changed files with 2 additions and 1 deletions
--- a/chat.py
+++ b/chat.py
@@ -93,7 +93,8 @@ class DOLPHIN:
            ckpt = try_rename_lagacy_weights(ckpt)
            self.model.load_state_dict(ckpt, strict=True)

-        self.model.to("cuda")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(device)
        self.model.eval()
        transform_args = {
            "input_size": self.swin_args["img_size"],