Update

[ghstack-poisoned]
pytorch · Dec 16, 2024 · 5721ce5 · 5721ce5
1 parent 6b4284f
commit 5721ce5
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/sota-implementations/td3/td3.py b/sota-implementations/td3/td3.py
@@ -38,6 +38,9 @@
 )
 
 
+torch.set_float32_matmul_precision("high")
+
+
 @hydra.main(version_base="1.1", config_path="", config_name="config")
 def main(cfg: "DictConfig"):  # noqa: F821
     device = cfg.network.device
@@ -196,9 +199,9 @@ def update(sampled_tensordict, update_actor):
                         torch.compiler.cudagraph_mark_step_begin()
                         q_loss, actor_loss = update(sampled_tensordict, update_actor)
 
-                    q_losses.append(q_loss)
+                    q_losses.append(q_loss.clone())
                     if update_actor:
-                        actor_losses.append(actor_loss)
+                        actor_losses.append(actor_loss.clone())
 
                     # Update priority
                     if prb: