Get params from task instead of model

Signed-off-by: Francesc Marti Escofet <f.martiescofet@gmail.com>

terratorch/tasks/base_task.py

@@ -58,11 +58,11 @@ def configure_optimizers(
         if self.hparams.get("lr_overrides", None) is not None and len(self.hparams["lr_overrides"]) > 0:
             parameters = []
             for param_name, custom_lr in self.hparams["lr_overrides"].items():
-                p = [p for n, p in self.model.named_parameters() if param_name in n]
+                p = [p for n, p in self.named_parameters() if param_name in n]
                 parameters.append({"params": p, "lr": custom_lr})
             rest_p = [
                 p
-                for n, p in self.model.named_parameters()
+                for n, p in self.named_parameters()
                 if all(param_name not in n for param_name in self.hparams["lr_overrides"])
             ]
             parameters.append({"params": rest_p})