diff --git a/tf-ner-poc/src/main/java/org/apache/opennlp/normalizer/Normalizer.java b/tf-ner-poc/src/main/java/org/apache/opennlp/normalizer/Normalizer.java
index c3f7fbb..9c9f27b 100644
--- a/tf-ner-poc/src/main/java/org/apache/opennlp/normalizer/Normalizer.java
+++ b/tf-ner-poc/src/main/java/org/apache/opennlp/normalizer/Normalizer.java
@@ -99,9 +99,8 @@
           .fetch("decode", 0).run();
 
       try (Tensor<?> translationTensor = result.get(0)) {
-        // TODO: This can't be hard coded ... normalized form doesn't need to have static length
         int[][] translations =
-            translationTensor.copyTo(new int[texts.length][9]); // shape is (20, 9) in eval py code
+            translationTensor.copyTo(new int[texts.length][(int) translationTensor.shape()[1]]);
 
         List<String> normalizedTexts = new ArrayList<>();