From 43305376d06c46101479290f4a55b4dddacb80f0 Mon Sep 17 00:00:00 2001 From: David Huggins-Daines Date: Tue, 16 Jul 2024 11:27:58 -0400 Subject: [PATCH] fix: WTF torch --- alexi/segment.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/alexi/segment.py b/alexi/segment.py index d372126..5be10a2 100644 --- 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)