diff --git a/cavis-zoo/cavis-zoo-models/src/main/java/org/deeplearning4j/zoo/model/TextGenerationLSTM.java b/cavis-zoo/cavis-zoo-models/src/main/java/org/deeplearning4j/zoo/model/TextGenerationLSTM.java
index 2a3475903..432c74231 100644
--- a/cavis-zoo/cavis-zoo-models/src/main/java/org/deeplearning4j/zoo/model/TextGenerationLSTM.java
+++ b/cavis-zoo/cavis-zoo-models/src/main/java/org/deeplearning4j/zoo/model/TextGenerationLSTM.java
@@ -47,7 +47,7 @@ public class TextGenerationLSTM extends ZooModel {
     @Builder.Default private long seed = 1234;
     @Builder.Default private int maxLength = 40;
     @Builder.Default private int totalUniqueCharacters = 47;
-    @Builder.Default private int[] inputShape = new int[] {maxLength, totalUniqueCharacters};
+    @Builder.Default private int[] inputShape = new int[] {40, 47};
     @Builder.Default private IUpdater updater = new RmsProp(0.01);
     @Builder.Default private CacheMode cacheMode = CacheMode.NONE;
     @Builder.Default private WorkspaceMode workspaceMode = WorkspaceMode.ENABLED;