diff --git a/pykoopman/koopman.py b/pykoopman/koopman.py
index 21ee2fb..d37d2e5 100644
--- a/pykoopman/koopman.py
+++ b/pykoopman/koopman.py
@@ -194,7 +194,7 @@ def simulate(self, x0, u=None, n_steps=1):
         y = empty((n_steps, self.n_input_features_), dtype=self.koopman_matrix.dtype)
         if u is None:
             y[0] = self.predict(x0)
-        elif u is not None:
+        else:
             y[0] = self.predict(x0, u[0])
 
         if isinstance(self.observables, TimeDelay):
@@ -207,10 +207,12 @@ def simulate(self, x0, u=None, n_steps=1):
         else:
             if u is None:
                 for k in range(n_steps - 1):
-                    y[k + 1] = self.predict(y[k])
+                    y[k + 1] = self.predict(y[k].reshape(1, -1))
             else:
                 for k in range(n_steps - 1):
-                    y[k + 1] = self.predict(y[k], u[k + 1])
+                    y[k + 1] = self.predict(
+                        y[k].reshape(1, -1), u[k + 1].reshape(1, -1)
+                    )
 
         return y