fix: WTF torch

dhdaines · Jul 16, 2024 · 4330537 · 4330537
1 parent 115e4ab
commit 4330537
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/alexi/segment.py b/alexi/segment.py
@@ -779,7 +779,7 @@ def __init__(self, model: PathLike = DEFAULT_RNN_MODEL, device="cpu"):
             self.model = RNNCRF(**self.config)
         else:
             self.model = RNN(**self.config)
-        self.model.load_state_dict(torch.load(model))
+        self.model.load_state_dict(torch.load(model, map_location=torch.device("cpu")))
         self.model.eval()
         self.model.to(device)