address #9

lucidrains · Dec 2, 2024 · ec10698 · ec10698
1 parent 9f733b0
commit ec10698
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/pi_zero_pytorch/pi_zero.py b/pi_zero_pytorch/pi_zero.py
@@ -1164,7 +1164,7 @@ def forward(
 
                 action_tokens = actions_ff(action_tokens) + action_tokens
 
-                action_tokens = ff_ada_rmsnorm(action_tokens, time_cond)
+                action_tokens = ff_ada_layerscale(action_tokens, time_cond)
 
         else:
 
@@ -1189,7 +1189,7 @@ def forward(
 
                 action_tokens = actions_ff(action_tokens) + action_tokens
 
-                action_tokens = ff_ada_rmsnorm(action_tokens, time_cond)
+                action_tokens = ff_ada_layerscale(action_tokens, time_cond)
 
         if not inferencing:
             # unpack and unembed to predictions

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "pi-zero-pytorch"
-version = "0.0.33"
+version = "0.0.34"
 description = "π0 in Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }