save encoder and decoder models independently

abcucberkeley · Jan 7, 2025 · 914cce3 · 914cce3
1 parent 76ca38d
commit 914cce3
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/training.py b/src/training.py
@@ -746,6 +746,8 @@ def trace_handler(prof):
                             accelerator.save_state(config['checkpointdir'] / f'best_state')
                             unwrapped_model = accelerator.unwrap_model(model)
                             torch.save(unwrapped_model.state_dict(), config['checkpointdir'] / f"best_model.bin")
+                            torch.save(unwrapped_model.get_encoder().state_dict(), config['checkpointdir'] / f"best_encoder.bin")
+                            torch.save(unwrapped_model.get_decoder().state_dict(), config['checkpointdir'] / f"best_decoder.bin")
 
                     checkpoint = Checkpoint.from_directory(checkpointdir)