diff --git a/src/main/java/ai/nets/samj/models/EfficientSamJ.java b/src/main/java/ai/nets/samj/models/EfficientSamJ.java
index b734654..cc1fe52 100644
--- a/src/main/java/ai/nets/samj/models/EfficientSamJ.java
+++ b/src/main/java/ai/nets/samj/models/EfficientSamJ.java
@@ -393,6 +393,7 @@ protected <T extends RealType<T> & NativeType<T>> void setImageOfInterest(Random
 		long[] dims = rai.dimensionsAsLongArray();
 		if (dims.length == 2)
 			rai = Views.addDimension(rai, 0, 0);
+		dims = rai.dimensionsAsLongArray();
 		if (dims[2] == 1)
 			rai = Views.interval( Views.expandMirrorDouble(rai, new long[] {0, 0, 2}), 
 					Intervals.createMinMax(new long[] {0, 0, 0, dims[0] - 1, dims[1] - 1, 2}) );
diff --git a/src/main/java/ai/nets/samj/models/EfficientViTSamJ.java b/src/main/java/ai/nets/samj/models/EfficientViTSamJ.java
index 5104ae1..6283e6a 100644
--- a/src/main/java/ai/nets/samj/models/EfficientViTSamJ.java
+++ b/src/main/java/ai/nets/samj/models/EfficientViTSamJ.java
@@ -303,6 +303,7 @@ protected <T extends RealType<T> & NativeType<T>> void setImageOfInterest(Random
 		long[] dims = rai.dimensionsAsLongArray();
 		if (dims.length == 2)
 			rai = Views.addDimension(rai, 0, 0);
+		dims = rai.dimensionsAsLongArray();
 		if (dims[2] == 1)
 			rai = Views.interval( Views.expandMirrorDouble(rai, new long[] {0, 0, 2}), 
 					Intervals.createMinMax(new long[] {0, 0, 0, dims[0] - 1, dims[1] - 1, 2}) );