From 5d2c3e2b4ac47b6bc34675aea9565aa3313c3b6c Mon Sep 17 00:00:00 2001 From: Filip Stefaniuk Date: Sat, 14 Sep 2024 04:32:42 -0400 Subject: [PATCH] add code for loading informer --- src/ml/model.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/ml/model.py b/src/ml/model.py index fee7a54..91bebbf 100644 --- 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")