roywei
diff --git a/‎tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/c_api/AbstractTF_Buffer.java‎
Lines changed: 9 additions & 3 deletions b/‎tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/c_api/AbstractTF_Buffer.java‎
Lines changed: 9 additions & 3 deletions
@@ -72,7 +72,8 @@ public static TF_Buffer newBufferFromString(Pointer proto) {
     }
 
     /**
-     * Returns a copy of the data in a Java array, or throws IndexOutOfBoundsException if too large.
+     * @return a copy of the data in a Java array
+     * @throws IndexOutOfBoundsException if too large.
      */
     public byte[] copyData() {
         long length = ((TF_Buffer)this).length();
@@ -85,10 +86,15 @@ public byte[] copyData() {
     }
 
     /**
-     * Returns the data of this buffer as a {@link java.nio.ByteBuffer}
+     * @return the data of this buffer as a {@link java.nio.ByteBuffer}
+     * @throws IndexOutOfBoundsException if too large.
      */
     public ByteBuffer dataAsByteBuffer() {
-        return ((TF_Buffer)this).data().capacity(((TF_Buffer) this).length()).asByteBuffer();
+        long length = ((TF_Buffer)this).length();
+        if (length > Integer.MAX_VALUE) {
+            throw new IndexOutOfBoundsException("TF_Buffer is too large to accessed via a ByteBuffer interface");
+        }
+        return ((TF_Buffer)this).data().capacity(length).asByteBuffer();
     }
 
     /**