Merge pull request #823 from marrlab/trainer_as_model

treat trainer as model
marrlab · May 7, 2024 · 4970eb6 · 4970eb6
2 parents ad94e79 + 79a2524
commit 4970eb6
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/domainlab/algos/trainers/a_trainer.py b/domainlab/algos/trainers/a_trainer.py
@@ -127,7 +127,9 @@ def init_business(self, model, task, observer, device, aconf, flag_accept=True):
             self._decoratee.init_business(
                 model, task, observer, device, aconf, flag_accept
             )
-        self.model = model
+            self.model = self._decoratee
+        else:
+            self.model = model
         self.task = task
         self.task.init_business(trainer=self, args=aconf)
         self.model.list_d_tr = self.task.list_domain_tr