hongshaorou/tdmpc2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use torch.compiler.cudagraph_mark_step_begin() in eval

Browse Source
This commit is contained in:
Vincent Moens
2024-10-26 00:32:16 +01:00
committed by GitHub
parent 836547d76f
commit fad0d1be03
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
tdmpc2/trainer/online_trainer.py
View File

@@ -31,6 +31,7 @@ class OnlineTrainer(Trainer):
if self.cfg.save_video: if self.cfg.save_video:
self.logger.video.init(self.env, enabled=(i==0)) self.logger.video.init(self.env, enabled=(i==0))
while not done: while not done:
torch.compiler.cudagraph_mark_step_begin()
action = self.agent.act(obs, t0=t==0, eval_mode=True) action = self.agent.act(obs, t0=t==0, eval_mode=True)
obs, reward, done, info = self.env.step(action) obs, reward, done, info = self.env.step(action)
ep_reward += reward ep_reward += reward