Fix setting model.criterio, using kwargs instead of hparams.

IGNF · Feb 1, 2024 · 25c6142 · 25c6142
1 parent 15f7bf6
commit 25c6142
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/myria3d/models/model.py b/myria3d/models/model.py
@@ -64,11 +64,11 @@ def __init__(self, **kwargs):
         # it also allows to access params with 'self.hparams' attribute
         self.save_hyperparameters(ignore=["criterion"])
 
-        neural_net_class = get_neural_net_class(self.hparams.neural_net_class_name)
-        self.model = neural_net_class(**self.hparams.neural_net_hparams)
+        neural_net_class = get_neural_net_class(kwargs.get("neural_net_class_name"))
+        self.model = neural_net_class(**kwargs.get("neural_net_hparams"))
 
         self.softmax = nn.Softmax(dim=1)
-        self.criterion = self.hparams.criterion
+        self.criterion = kwargs.get("criterion")
 
     def on_fit_start(self) -> None:
         self.criterion = self.criterion.to(self.device)