Fix Metaphor Annotation for use without GPU.

cvjena · Oct 26, 2024 · fdec8b3 · fdec8b3
1 parent db32bf5
commit fdec8b3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/setup.cfg b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = freestylo
-version = 0.2.0
+version = 0.3.0
 description = A tool for stylistic device detection."
 author = Felix Schneider
 author_email = [email protected]

diff --git a/src/freestylo/MetaphorAnnotation.py b/src/freestylo/MetaphorAnnotation.py
@@ -36,7 +36,7 @@ def serialize(self) -> list:
     def load_model(self, model_path):
         model_path = get_model_path(model_path)
         self.model = SimilarityNN.SimilarityNN(300, 128, 1, 128, self.device)
-        self.model.load_state_dict(torch.load(model_path, weights_only=True))
+        self.model.load_state_dict(torch.load(model_path, weights_only=True, map_location=self.device))
         self.model = self.model.to(self.device)
         self.model.eval()