amend

pytorch · Apr 8, 2024 · 980a7ae · 980a7ae
1 parent 07d3e93
commit 980a7ae
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sota-implementations/dreamer/dreamer.py b/sota-implementations/dreamer/dreamer.py
@@ -129,8 +129,8 @@ def main(cfg: "DictConfig"):  # noqa: F821
     def compile_rssms(module):
         if isinstance(module, RSSMRollout) and not getattr(module, "_compiled", False):
             module._compiled = True
-            module.rssm_prior.module = torch.compile(module.rssm_prior.module)
-            module.rssm_posterior.module = torch.compile(module.rssm_posterior.module)
+            module.rssm_prior.module = torch.compile(module.rssm_prior.module, backend="cudagraphs")
+            module.rssm_posterior.module = torch.compile(module.rssm_posterior.module, backend="cudagraphs")
     world_model_loss.apply(compile_rssms)
 
     t_collect_init = time.time()