Fix data setter for BNN targets

Author: ziatdinovmax

gpax/models/bnn.py

@@ -32,7 +32,7 @@ def _set_data(self, X: jnp.ndarray, y: Optional[jnp.ndarray] = None
                   ) -> Union[Tuple[jnp.ndarray], jnp.ndarray]:
         X = X if X.ndim > 1 else X[:, None]
         if y is not None:
-            y = y[:, None] if y.ndim < 1 else y
+            y = y[:, None] if y.ndim < 2 else y
             return X, y
         return X