dnrops
/
transformers
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add Sourab's comments

This commit is contained in:
Zach Mueller 2024-05-15 14:50:49 -04:00
parent a39094919d
commit 1a5009640b
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/transformers/trainer.py
View File

@ -4591,7 +4591,10 @@ class Trainer:
>>> model = trainer.free_memory(model)
```
"""
*models, self.optimizer = release_memory(*models, self.optimizer)
# We need to have these references so they can be set to `None`
*models, self.optimizer, self.model, self.deepspeed, self.model_wrapped = release_memory(
*models, self.optimizer, self.model, self.deepspeed, self.model_wrapped
)
return models
def propagate_args_to_deepspeed(self, auto_find_batch_size=False):