Fix issue with MLP

Ceyron · Mar 12, 2024 · f46be98 · f46be98
1 parent 99c081e
commit f46be98
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pdequinox/_mlp.py b/pdequinox/_mlp.py
@@ -57,11 +57,11 @@ def __init__(
         )
 
     def __call__(self, x):
-        if x.shape != self.in_shape:
+        if x.shape != self._in_shape:
             raise ValueError(
-                f"Input shape {x.shape} does not match expected shape {self.in_shape}. For batched operation use jax.vmap"
+                f"Input shape {x.shape} does not match expected shape {self._in_shape}. For batched operation use jax.vmap"
             )
         x_flat = x.flatten()
         x_flat = self.flat_mlp(x_flat)
-        x = x_flat.reshape(self.out_shape)
+        x = x_flat.reshape(self._out_shape)
         return x