Update

[ghstack-poisoned]
pytorch · Dec 16, 2024 · 6b4284f · 6b4284f
1 parent 2b39661
commit 6b4284f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sota-implementations/td3/td3.py b/sota-implementations/td3/td3.py
@@ -221,9 +221,9 @@ def update(sampled_tensordict, update_actor):
             )
 
         if collected_frames >= init_random_frames:
-            metrics_to_log["train/q_loss"] = q_losses.mean()
+            metrics_to_log["train/q_loss"] = torch.stack(q_losses).mean()
             if update_actor:
-                metrics_to_log["train/a_loss"] = actor_losses.mean()
+                metrics_to_log["train/a_loss"] = torch.stack(actor_losses).mean()
 
         # Evaluation
         if abs(collected_frames % eval_iter) < frames_per_batch: