Update

[ghstack-poisoned]
pytorch · Dec 15, 2024 · a4abc7b · a4abc7b
1 parent ad97990
commit a4abc7b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/sota-implementations/gail/gail.py b/sota-implementations/gail/gail.py
@@ -34,6 +34,9 @@
 from torchrl.record.loggers import generate_exp_name, get_logger
 
 
+torch.set_float32_matmul_precision("high")
+
+
 @hydra.main(config_path="", config_name="config")
 def main(cfg: "DictConfig"):  # noqa: F821
     set_gym_backend(cfg.env.backend).set()
@@ -260,7 +263,7 @@ def update(data, expert_data, num_network_updates=num_network_updates):
         expert_data = expert_data.to(device)
 
         metadata = update(data, expert_data)
-        d_loss = metadata["d_loss"]
+        d_loss = metadata["dloss"]
         alpha = metadata["alpha"]
 
         # Get training rewards and episode lengths

diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -3319,9 +3319,9 @@ def __init__(
         self.update_mask(mask)
         self._provisional_n = None
 
-    @torch.compiler.assume_constant_result
+    @property
     def _undefined_n(self):
-        return self.space.n == -1
+        return self.space.n < 0
 
     def enumerate(self) -> torch.Tensor:
         dtype = self.dtype