From 377d3d189eea1a068d65f8918a2b7fbc7d1a1977 Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Thu, 14 Jul 2022 16:21:31 +0800
Subject: Make class variables available for `_custom_init_fn`

---
 libs/utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

(limited to 'libs')

diff --git a/libs/utils.py b/libs/utils.py
index bc45a12..5e8529b 100644
--- a/libs/utils.py
+++ b/libs/utils.py
@@ -93,8 +93,6 @@ class Trainer(ABC):
             optims.items(), last_iter, num_iters, config.sched_config,
         ))
 
-        self._custom_init_fn(config)
-
         self.restore_iter = last_iter + 1
         self.train_loader = train_loader
         self.test_loader = test_loader
@@ -104,6 +102,8 @@ class Trainer(ABC):
         self._inf_mode = inf_mode
         self._checkpoint_dir = checkpoint_dir
 
+        self._custom_init_fn(config)
+
     @dataclass
     class BatchLogRecord(BaseBatchLogRecord):
         pass
-- 
cgit v1.2.3