From 79a2524560f8a654ff4032dc992d35f4b7183027 Mon Sep 17 00:00:00 2001
From: smilesun <smilesun.east@gmail.com>
Date: Tue, 7 May 2024 16:39:28 +0200
Subject: [PATCH] .

---
 domainlab/algos/trainers/a_trainer.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/domainlab/algos/trainers/a_trainer.py b/domainlab/algos/trainers/a_trainer.py
index 9bc126705..d9bda1513 100644
--- 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