fix loss

Acellera · Apr 16, 2024 · 96effb1 · 96effb1
1 parent d5848ee
commit 96effb1
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/scripts/a2c/a2c.py b/scripts/a2c/a2c.py
@@ -319,7 +319,9 @@ def create_env_fn():
             loss = loss_module(batch)
             loss = loss.named_apply(
                 lambda name, value: (
-                    (value * mask).mean() if name.startswith("loss_") else value
+                    (value * mask).mean()
+                    if name.startswith("loss_")
+                    else value
                     # (value * mask).sum(-1).mean(-1) if name.startswith("loss_") else value
                 ),
                 batch_size=[],

diff --git a/scripts/ppo/ppo.py b/scripts/ppo/ppo.py
@@ -360,7 +360,9 @@ def create_env_fn():
                 loss = loss_module(batch)
                 loss = loss.named_apply(
                     lambda name, value: (
-                        value * mask).mean() if name.startswith("loss_") else value
+                        (value * mask).mean()
+                        if name.startswith("loss_")
+                        else value
                         # (value * mask).sum(-1).mean(-1) if name.startswith("loss_") else value
                     ),
                     batch_size=[],