Fix.

e3nn · Sep 28, 2024 · 245e17e · 245e17e
1 parent aa15f91
commit 245e17e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/e3nn_jax/_src/linear_flax.py b/e3nn_jax/_src/linear_flax.py
@@ -74,9 +74,9 @@ class Linear(flax.linen.Module):
     gradient_normalization: Optional[Union[float, str]] = None
     path_normalization: Optional[Union[float, str]] = None
     biases: bool = False
-    parameter_initializer: Optional[
-        Callable[[], jax.nn.initializers.Initializer]
-    ] = None
+    parameter_initializer: Optional[Callable[[], jax.nn.initializers.Initializer]] = (
+        None
+    )
     instructions: Optional[List[Tuple[int, int]]] = None
     num_indexed_weights: Optional[int] = None
     weights_per_channel: bool = False