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

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

Browse Source
This commit is contained in:
Filip Stefaniuk 2024-09-14 10:32:56 +02:00
commit 747c5cda43
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")