Merge branch 'main' of gitlab.com:FilipStefaniuk/wne-msc-thesis

2024-09-14 10:32:56 +02:00 · 2024-09-14 10:32:56 +02:00 · 747c5cda43
commit 747c5cda43
parent 346c3011e8 5d2c3e2b4a
1 changed files with 2 additions and 0 deletions
--- a/src/ml/model.py
+++ b/src/ml/model.py
@ -72,6 +72,8 @@ def load_model_from_wandb(run):
    if model_name == 'TemporalFusionTransformer':
        return TemporalFusionTransformer.load_from_checkpoint(
            model_artifact.file())
+    if model_name == 'Informer':
+        return Informer.load_from_checkpoint(model_artifact.file())

    raise ValueError("Invalid model name")