test loading using generic model

webis-de · Sep 4, 2024 · fdd49df · fdd49df
1 parent 080ff18
commit fdd49df
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/test_model.py b/tests/test_model.py
@@ -4,7 +4,7 @@
 import pytest
 from transformers import AutoModel
 
-from lightning_ir.base import LightningIRModule
+from lightning_ir.base import LightningIRModel, LightningIRModule
 from lightning_ir.cross_encoder import CrossEncoderModule
 from lightning_ir.data import LightningIRDataModule, RunDataset, TupleDataset
 from lightning_ir.loss.loss import InBatchLossFunction
@@ -68,6 +68,7 @@ def test_seralize_deserialize(module: LightningIRModule, tmp_path: Path):
     new_models = [
         model.__class__.from_pretrained(save_dir),
         model.__class__.__bases__[0].from_pretrained(save_dir),
+        LightningIRModel.from_pretrained(save_dir),
         AutoModel.from_pretrained(save_dir),
     ]
     for new_model in new_models: