Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
update
Browse files Browse the repository at this point in the history
songhappy committed Jan 23, 2025

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
1 parent aa836ec commit da30adb
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions recipes/ppo_full_finetune_single_device.py
Original file line number Diff line number Diff line change
@@ -935,6 +935,7 @@ def train(self) -> None:
curr_epoch == 0
and self.profiler_profile_memory
and idx == self.profiler_wait_steps + self.profiler_warmup_steps
and self._device.type == "cuda"
):
torch.cuda.memory._record_memory_history()

@@ -1034,6 +1035,7 @@ def train(self) -> None:
== self.profiler_wait_steps
+ self.profiler_warmup_steps
+ self.profiler_active_steps
and self._device.type == "cuda"
):
torch.cuda.memory._record_memory_history(enabled=None)

0 comments on commit da30adb

Please sign in to comment.