mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

i like wasting time

Browse Source
This commit is contained in:
James Betker 2022-01-24 18:12:08 -07:00
parent fc09cff4b3
commit 798ed7730a
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
codes/train.py
View File

@ -215,7 +215,7 @@ class Trainer:
#### save models and training states
if self.current_step % opt['logger']['save_checkpoint_freq'] == 0:
self.model.consolidate_state(state)
self.model.consolidate_state()
if self.rank <= 0:
self.logger.info('Saving models and training states.')
self.model.save(self.current_step)

4
codes/trainer/base_model.py
View File

@ -129,10 +129,10 @@ class BaseModel():
network.load_state_dict(load_net_clean, strict=strict)
def consolidate_state(self, state):
def consolidate_state(self):
for o in self.optimizers:
if isinstance(o, ZeroRedundancyOptimizer):
state['optimizers'].append(o.consolidate_state_dict(to=0))
o.consolidate_state_dict(to=0)
def save_training_state(self, state):