add weatherbench to the architectures pool

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

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()