diff --git a/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/autodiff/samediff/internal/memory/ArrayCacheMemoryMgr.java b/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/autodiff/samediff/internal/memory/ArrayCacheMemoryMgr.java
index 6fb12b826..69db24f03 100644
--- a/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/autodiff/samediff/internal/memory/ArrayCacheMemoryMgr.java
+++ b/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/autodiff/samediff/internal/memory/ArrayCacheMemoryMgr.java
@@ -252,11 +252,13 @@ public class ArrayCacheMemoryMgr extends AbstractMemoryMgr {
         private void removeObject(INDArray array){
             long length = array.data().length();
             int idx = Arrays.binarySearch(lengths, 0, size, length);
-            Preconditions.checkState(idx > 0, "Cannot remove array from ArrayStore: no array with this length exists in the cache");
+            Preconditions.checkState(idx >= 0,
+                    "Cannot remove array from ArrayStore: no array with this length exists in the cache");
             boolean found = false;
             int i = 0;
-            while(!found && i <= size && lengths[i] == length){
-                found = sorted[i++] == array; //Object equality
+            while (!found && i < size) {
+                found = sorted[i] == array && lengths[i] == length; //Object and length equality
+                ++i;
             }
             Preconditions.checkState(found, "Cannot remove array: not found in ArrayCache");
             removeIdx(i - 1);