removed scalar activation

atomicarchitects · Nov 24, 2024 · 2abe2c5 · 2abe2c5
1 parent 9172669
commit 2abe2c5
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/tetris_jax_baseline.py b/tetris_jax_baseline.py
@@ -68,8 +68,6 @@ def update_edge_fn(edge_features, sender_features, receiver_features, globals):
         def update_node_fn(node_features, sender_features, receiver_features, globals):
             node_feats = receiver_features / self.denominator
             node_feats = e3nn.flax.Linear(target_irreps, name="linear_pre")(node_feats)
-            node_feats = e3nn.scalar_activation(node_feats)
-            node_feats = e3nn.flax.Linear(target_irreps, name="linear_post")(node_feats)
             shortcut = e3nn.flax.Linear(
                 node_feats.irreps, name="shortcut", force_irreps_out=True
             )(node_features)