🐛 Bug
To Reproduce
import torch
from litmodels import save_model, load_model
model = torch.nn.Module()
save_model(model=model, name="skaftenicki/language-model/torch-model")
model_ = load_model(name="skaftenicki/language-model/torch-model")
generates
❯ python main.py
Uploading files...: 100%|█████████████████| 1/1 [00:01<00:00, 1.73s/it]
Model uploaded successfully. Link to the model: 'https://lightning.ai/skaftenicki/language-model/models/torch-model'
Downloading default: 100%|█████████| 1.32k/1.32k [00:00<00:00, 1.71kB/s]
Traceback (most recent call last): | 0.00/1.32k [00:00<?, ?B/s]
File "/home/nsde/Documents/LitModels/main.py", line 6, in <module>
model_ = load_model(name="skaftenicki/language-model/torch-model")
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/nsde/Documents/LitModels/src/litmodels/io/gateway.py", line 166, in load_model
raise NotImplementedError(f"Loading model from {model_path.suffix} is not supported yet.")
NotImplementedError: Loading model from .pth is not supported yet.
Expected behavior
Code is basically taken from the examples. The problem seems to be that save_model function saves torch models to pth (which are not ScriptModule):
|
elif _PYTORCH_AVAILABLE and isinstance(model, torch.nn.Module): |
|
path = os.path.join(staging_dir, f"{model.__class__.__name__}.pth") |
|
torch.save(model.state_dict(), path) |
but not supported in the
load_model function.
Environment
- PyTorch Version (e.g., 1.0):
- OS (e.g., Linux):
- How you installed PyTorch (
conda, pip, source):
- Build command you used (if compiling from source):
- Python version:
- CUDA/cuDNN version:
- GPU models and configuration:
- Any other relevant information:
Additional context
🐛 Bug
To Reproduce
generates
Expected behavior
Code is basically taken from the examples. The problem seems to be that
save_modelfunction saves torch models topth(which are notScriptModule):LitModels/src/litmodels/io/gateway.py
Lines 94 to 96 in e9075da
but not supported in the
load_modelfunction.Environment
conda,pip, source):Additional context