From f93d049759f7920ff49d8a195bab3a15edee8b06 Mon Sep 17 00:00:00 2001 From: deependujha Date: Fri, 25 Sep 2026 13:43:04 +0530 Subject: [PATCH] fix: support loading .pth models --- src/litmodels/io/gateway.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/litmodels/io/gateway.py b/src/litmodels/io/gateway.py index b0d7d2d..afc73a5 100644 --- a/src/litmodels/io/gateway.py +++ b/src/litmodels/io/gateway.py @@ -155,6 +155,7 @@ def load_model(name: str, download_dir: str = ".") -> Any: Supported formats: - .ts → torch.jit.load + - .pth → torch.load - .keras → keras.models.load_model - .pkl → pickle/joblib via load_pickle @@ -180,4 +181,6 @@ def load_model(name: str, download_dir: str = ".") -> Any: return keras.models.load_model(model_path) if model_path.suffix.lower() == ".pkl": return load_pickle(path=model_path) + if model_path.suffix.lower() == ".pth": + return torch.load(model_path, weights_only=False) raise NotImplementedError(f"Loading model from {model_path.suffix} is not supported yet.")