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

update compile print

Browse Source
This commit is contained in:
Nicklas Hansen
2024-11-10 12:27:24 -08:00
parent c694d286f0
commit fb07cdac3f
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tdmpc2/tdmpc2.py
View File

@@ -36,7 +36,7 @@ class TDMPC2(torch.nn.Module):
) if self.cfg.multitask else self._get_discount(cfg.episode_length) ) if self.cfg.multitask else self._get_discount(cfg.episode_length)
self._prev_mean = torch.nn.Buffer(torch.zeros(self.cfg.horizon, self.cfg.action_dim, device=self.device)) self._prev_mean = torch.nn.Buffer(torch.zeros(self.cfg.horizon, self.cfg.action_dim, device=self.device))
if cfg.compile: if cfg.compile:
print('compiling - update') print('Compiling update function with torch.compile...')
self._update = torch.compile(self._update, mode="reduce-overhead") self._update = torch.compile(self._update, mode="reduce-overhead")
@property @property