From f098224b8a553d00f07a167440ca5c3acb698aeb Mon Sep 17 00:00:00 2001
From: BING GONG <b.gong@fz-juelich.de>
Date: Tue, 7 Jun 2022 15:45:25 +0200
Subject: [PATCH] add weatherbench to the architectures pool
---
video_prediction_tools/model_modules/model_architectures.py | 1 +
.../model_modules/video_prediction/models/__init__.py | 3 +--
2 files changed, 2 insertions(+), 2 deletions(-)
diff --git a/video_prediction_tools/model_modules/model_architectures.py b/video_prediction_tools/model_modules/model_architectures.py
index b8793ad8..b33ed8c5 100644
--- a/video_prediction_tools/model_modules/model_architectures.py
+++ b/video_prediction_tools/model_modules/model_architectures.py
@@ -13,6 +13,7 @@ def known_models():
'convLSTM_gan': "ConvLstmGANVideoPredictionModel",
'ours_vae_l1': 'SAVPVideoPredictionModel',
'ours_gan': 'SAVPVideoPredictionModel',
+ "weatherBench": "WeatherBenchModel "
}
return model_mappings
diff --git a/video_prediction_tools/model_modules/video_prediction/models/__init__.py b/video_prediction_tools/model_modules/video_prediction/models/__init__.py
index a4b60965..290def9f 100644
--- a/video_prediction_tools/model_modules/video_prediction/models/__init__.py
+++ b/video_prediction_tools/model_modules/video_prediction/models/__init__.py
@@ -14,8 +14,7 @@ from .mcnet_model import McNetVideoPredictionModel
from .test_model import TestModelVideoPredictionModel
from model_modules.model_architectures import known_models
from .convLSTM_GAN_model import ConvLstmGANVideoPredictionModel
-
-
+from .weatherBench3DCNN import WeatherBenchModel
def get_model_class(model):
model_mappings = known_models()
--
GitLab