Merge pull request #31 from nasaharvest/model-device

Move presto pretrained model to the right device
nasaharvest · Dec 23, 2023 · 632cd39 · 632cd39
2 parents f6422ad + 6d8c6c9
commit 632cd39
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/presto/presto.py b/presto/presto.py
@@ -789,4 +789,4 @@ def construct_finetuning_model(
     def load_pretrained(cls, model_path: Union[str, Path] = default_model_path):
         model = cls.construct()
         model.load_state_dict(torch.load(model_path, map_location=device))
-        return model
+        return model.to(device)