yasha/informer_model
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add code for loading informer

Browse Source
This commit is contained in:
Filip Stefaniuk 2024-09-14 04:32:42 -04:00
parent 73c513c217
commit 5d2c3e2b4a
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/ml/model.py
View File

@ -72,6 +72,8 @@ def load_model_from_wandb(run):
if model_name == 'TemporalFusionTransformer': if model_name == 'TemporalFusionTransformer':
return TemporalFusionTransformer.load_from_checkpoint( return TemporalFusionTransformer.load_from_checkpoint(
model_artifact.file()) model_artifact.file())
if model_name == 'Informer':
return Informer.load_from_checkpoint(model_artifact.file())
raise ValueError("Invalid model name") raise ValueError("Invalid model name")