Freeze the traced PyTorch model

2025-07-11 00:02:14 +08:00 · 2022-06-21 13:43:18 +01:00 · 2022-06-21 13:43:18 +01:00 · 081c8efaa0
commit 081c8efaa0
parent ec041ffdfd
1 changed files with 1 additions and 0 deletions
--- a/eland/ml/pytorch/traceable_model.py
+++ b/eland/ml/pytorch/traceable_model.py
@ -59,5 +59,6 @@ class TraceableModel(ABC):
    def save(self, path: str) -> str:
        model_path = os.path.join(path, "traced_pytorch_model.pt")
        trace_model = self.trace()
+        trace_model = torch.jit.freeze(trace_model)
        torch.jit.save(trace_model, model_path)
        return model_path