third_codes/eland
1
0
Fork 0
mirror of https://github.com/elastic/eland.git synced 2025-07-11 00:02:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Freeze the traced PyTorch model

Browse Source
This commit is contained in:
David Kyle 2022-06-21 13:43:18 +01:00 committed by GitHub
parent ec041ffdfd
commit 081c8efaa0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
eland/ml/pytorch/traceable_model.py
View File

@ -59,5 +59,6 @@ class TraceableModel(ABC):
def save(self, path: str) -> str: def save(self, path: str) -> str:
model_path = os.path.join(path, "traced_pytorch_model.pt") model_path = os.path.join(path, "traced_pytorch_model.pt")
trace_model = self.trace() trace_model = self.trace()
trace_model = torch.jit.freeze(trace_model)
torch.jit.save(trace_model, model_path) torch.jit.save(trace_model, model_path)
return model_path return model_path