diff --git a/README.md b/README.md
index ace45543c..a31d8348c 100644
--- a/README.md
+++ b/README.md
@@ -46,8 +46,8 @@ X_train, X_test, y_train, y_test = \
         sklearn.model_selection.train_test_split(X, y, random_state=1)
 
 # running Auto-PyTorch
-autoPyTorch = AutoNetClassification(log_level='info')
-autoPyTorch.fit(X_train, y_train)
+autoPyTorch = AutoNetClassification(log_level='info', max_runtime=300, min_budget=30, max_budget=90)
+autoPyTorch.fit(X_train, y_train, validation_split=0.3)
 y_pred = autoPyTorch.predict(X_test)
 
 print("Accuracy score", sklearn.metrics.accuracy_score(y_test, y_pred))
diff --git a/autoPyTorch/core/api.py b/autoPyTorch/core/api.py
index 15e6b5563..a32641001 100644
--- a/autoPyTorch/core/api.py
+++ b/autoPyTorch/core/api.py
@@ -166,7 +166,7 @@ def predict(self, X, return_probabilities=False):
 
         # reverse one hot encoding 
         OHE = self.pipeline[OneHotEncoding.get_name()]
-        result = OHE.reverse_transform_y(Y_pred, OHE.fit_output['y_one_hot_encoder']).reshape(1, -1)
+        result = OHE.reverse_transform_y(Y_pred, OHE.fit_output['y_one_hot_encoder'])
         return result if not return_probabilities else (result, Y_pred)
 
     def score(self, X_test, Y_test):