Add LSTM type extraction

JonathanCrabbe · Jan 25, 2024 · f3456f1 · f3456f1
1 parent 6a85234
commit f3456f1
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/fdiff/utils/extraction.py b/src/fdiff/utils/extraction.py
@@ -6,7 +6,7 @@
 from omegaconf import DictConfig, OmegaConf
 
 from fdiff.dataloaders.datamodules import Datamodule
-from fdiff.models.score_models import MLPScoreModule, ScoreModule
+from fdiff.models.score_models import LSTMScoreModule, MLPScoreModule, ScoreModule
 
 
 def get_training_params(datamodule: Datamodule, trainer: pl.Trainer) -> dict[str, Any]:
@@ -70,6 +70,8 @@ def get_model_typle(cfg: DictConfig | dict) -> ScoreModule | MLPScoreModule:
             return ScoreModule
         case "fdiff.models.score_models.MLPScoreModule":
             return MLPScoreModule
+        case "fdiff.models.score_models.LSTMScoreModule":
+            return LSTMScoreModule
         case _:
             raise NotImplementedError(f"Model class {model_class} not implemented yet.")