diff --git a/tensorflow_lite_support/java/src/java/org/tensorflow/lite/support/common/ops/NormalizeOp.java b/tensorflow_lite_support/java/src/java/org/tensorflow/lite/support/common/ops/NormalizeOp.java
index 8ac57eed2..2097cf3e6 100644
--- a/tensorflow_lite_support/java/src/java/org/tensorflow/lite/support/common/ops/NormalizeOp.java
+++ b/tensorflow_lite_support/java/src/java/org/tensorflow/lite/support/common/ops/NormalizeOp.java
@@ -141,11 +141,11 @@ public TensorBuffer apply(@NonNull TensorBuffer input) {
     SupportPreconditions.checkArgument(
         numChannels == 1 || (shape.length != 0 && shape[shape.length - 1] == numChannels),
         "Number of means (stddevs) is not same with number of channels (size of last axis).");
-    // TODO(136750944): Eliminate the array copy here.
-    float[] values = input.getFloatArray();
+    int flatSize = input.getFlatSize();    
+    float[] values = new float[flatSize];
     int j = 0;
-    for (int i = 0; i < values.length; i++) {
-      values[i] = (values[i] - mean[j]) / stddev[j];
+    for (int i = 0; i < flatSize; i++) {
+      values[i] = (input.getFloatValue(i) - mean[j]) / stddev[j];
       j = (j + 1) % numChannels;
     }
     TensorBuffer output;