Respair
/

OTA_TextAligner

Model card Files Files and versions

Respair commited on Mar 30

Commit

faa8827

·

verified ·

1 Parent(s): 9bf7f12

Update ddp_train.py

Files changed (1) hide show

ddp_train.py +1 -1

ddp_train.py CHANGED Viewed

@@ -219,7 +219,7 @@ def main(config_path):
           print(f'Saving on step {epoch*len(train_dataloader)+i}...')
           state = {
-              'net':  {key: aligner[key].state_dict() for key in model},
               'optimizer': optimizer.state_dict(),
               'iters': iters,
               'epoch': epoch,

           print(f'Saving on step {epoch*len(train_dataloader)+i}...')
           state = {
+              'net':  {key: aligner[key].state_dict() for key in aligner},
               'optimizer': optimizer.state_dict(),
               'iters': iters,
               'epoch': epoch,