Skip to content
This repository has been archived by the owner on Nov 22, 2022. It is now read-only.

Commit

Permalink
remove move_state_dict_to_gpu, which is causing cuda oom (#1367)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: #1367

I keep getting cuda oom in this load_best_model stage, move_state_dict_to_gpu and model.cuda() are not both needed. Looks like it will double gpu memory this way.

Reviewed By: anchit

Differential Revision: D21725316

fbshipit-source-id: 70b5761a25afb19da7f44a3fead37b36d0e122da
  • Loading branch information
Haoran Li authored and facebook-github-bot committed May 27, 2020
1 parent 2fec533 commit b0a9d80
Showing 1 changed file with 1 addition and 3 deletions.
4 changes: 1 addition & 3 deletions pytext/trainers/trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -333,9 +333,7 @@ def load_best_model(self, state: TrainingState):
if cuda.CUDA_ENABLED:
# Move current model to CPU to avoid multiple models in GPU memory
state.model.cpu()
state.model.load_state_dict(
self.move_state_dict_to_gpu(state.best_model_state)
)
state.model.load_state_dict(state.best_model_state)
# Move model back to GPU
state.model.cuda()
else:
Expand Down

0 comments on commit b0a9d80

Please sign in to comment.