Update

[ghstack-poisoned]
pytorch · Dec 15, 2024 · ab3616c · ab3616c
1 parent c97088f
commit ab3616c
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/sota-implementations/iql/discrete_iql.py b/sota-implementations/iql/discrete_iql.py
@@ -129,11 +129,10 @@ def update(sampled_tensordict):
 
         # update qnet_target params
         target_net_updater.step()
-        return TensorDict(
-            metadata.update(
-                {"actor_loss": actor_loss, "value_loss": value_loss, "q_loss": q_loss}
-            )
-        ).detach()
+        metadata.update(
+            {"actor_loss": actor_loss, "value_loss": value_loss, "q_loss": q_loss}
+        )
+        return TensorDict(metadata).detach()
 
     if cfg.compile.compile:
         update = torch.compile(update, mode=compile_mode)