amend

pytorch · Apr 8, 2024 · 7c59b96 · 7c59b96
1 parent 0ee83cd
commit 7c59b96
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/sota-implementations/dreamer/dreamer.py b/sota-implementations/dreamer/dreamer.py
@@ -127,7 +127,8 @@ def main(cfg: "DictConfig"):  # noqa: F821
 
     print('Compiling')
     def compile_rssms(module):
-        if isinstance(module, RSSMRollout):
+        if isinstance(module, RSSMRollout) and not getattr(module, "_compiled", False):
+            module._compiled = True
             module.rssm_prior = torch.compile(module.rssm_prior)
             module.rssm_posterior = torch.compile(module.rssm_posterior)
     world_model_loss.apply(compile_rssms)