feat: make model instantiateable

ecmwf · Sep 5, 2024 · 813d36d · 813d36d
1 parent 5525e7a
commit 813d36d
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/src/anemoi/models/interface/__init__.py b/src/anemoi/models/interface/__init__.py
@@ -14,7 +14,6 @@
 from hydra.utils import instantiate
 from torch_geometric.data import HeteroData
 
-from anemoi.models.models.encoder_processor_decoder import AnemoiModelEncProcDec
 from anemoi.models.preprocessing import Processors
 
 
@@ -73,9 +72,9 @@ def _build_model(self) -> None:
         self.pre_processors = Processors(processors)
         self.post_processors = Processors(processors, inverse=True)
 
-        # Instantiate the model (Can be generalised to other models in the future, here we use AnemoiModelEncProcDec)
-        self.model = AnemoiModelEncProcDec(
-            config=self.config, data_indices=self.data_indices, graph_data=self.graph_data
+        # Instantiate the model
+        self.model = instantiate(
+            self.config.model.model, config=self.config, data_indices=self.data_indices, graph_data=self.graph_data
         )
 
         # Use the forward method of the model directly