correct error creating tensors

bioimage-io · Oct 27, 2023 · 5fcda4d · 5fcda4d
1 parent 6dc9d43
commit 5fcda4d
Showing 1 changed file with 10 additions and 10 deletions.
diff --git a/src/main/java/io/bioimage/modelrunner/tensorflow/v2/api020/tensor/TensorBuilder.java b/src/main/java/io/bioimage/modelrunner/tensorflow/v2/api020/tensor/TensorBuilder.java
@@ -131,6 +131,7 @@ private static Tensor<TUint8> buildUByte(
 		RandomAccessibleInterval<UnsignedByteType> tensor)
 		throws IllegalArgumentException
 	{
+		long[] ogShape = tensor.dimensionsAsLongArray();
 		tensor = Utils.transpose(tensor);
 		PrimitiveBlocks< UnsignedByteType > blocks = PrimitiveBlocks.of( tensor );
 		long[] tensorShape = tensor.dimensionsAsLongArray();
@@ -142,8 +143,7 @@ private static Tensor<TUint8> buildUByte(
 			sArr[i] = (int) tensorShape[i];
 		blocks.copy( new long[tensorShape.length], flatArr, sArr );
 		ByteDataBuffer dataBuffer = RawDataBufferFactory.create(flatArr, false);
-		Tensor<TUint8> ndarray = Tensor.of(TUint8.DTYPE, Shape.of(tensor
-			.dimensionsAsLongArray()), dataBuffer);
+		Tensor<TUint8> ndarray = Tensor.of(TUint8.DTYPE, Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}
 
@@ -160,6 +160,7 @@ private static Tensor<TUint8> buildUByte(
 	private static Tensor<TInt32> buildInt(
 		RandomAccessibleInterval<IntType> tensor) throws IllegalArgumentException
 	{
+		long[] ogShape = tensor.dimensionsAsLongArray();
 		tensor = Utils.transpose(tensor);
 		PrimitiveBlocks< IntType > blocks = PrimitiveBlocks.of( tensor );
 		long[] tensorShape = tensor.dimensionsAsLongArray();
@@ -171,8 +172,7 @@ private static Tensor<TInt32> buildInt(
 			sArr[i] = (int) tensorShape[i];
 		blocks.copy( new long[tensorShape.length], flatArr, sArr );
 		IntDataBuffer dataBuffer = RawDataBufferFactory.create(flatArr, false);
-		Tensor<TInt32> ndarray = TInt32.tensorOf(Shape.of(tensor
-			.dimensionsAsLongArray()), dataBuffer);
+		Tensor<TInt32> ndarray = TInt32.tensorOf(Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}
 
@@ -190,6 +190,7 @@ private static Tensor<TInt64> buildLong(
 		RandomAccessibleInterval<LongType> tensor)
 		throws IllegalArgumentException
 	{
+		long[] ogShape = tensor.dimensionsAsLongArray();
 		tensor = Utils.transpose(tensor);
 		PrimitiveBlocks< LongType > blocks = PrimitiveBlocks.of( tensor );
 		long[] tensorShape = tensor.dimensionsAsLongArray();
@@ -201,8 +202,7 @@ private static Tensor<TInt64> buildLong(
 			sArr[i] = (int) tensorShape[i];
 		blocks.copy( new long[tensorShape.length], flatArr, sArr );
 		LongDataBuffer dataBuffer = RawDataBufferFactory.create(flatArr, false);
-		Tensor<TInt64> ndarray = TInt64.tensorOf(Shape.of(tensor
-			.dimensionsAsLongArray()), dataBuffer);
+		Tensor<TInt64> ndarray = TInt64.tensorOf(Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}
 
@@ -220,6 +220,7 @@ private static Tensor<TFloat32> buildFloat(
 		RandomAccessibleInterval<FloatType> tensor)
 		throws IllegalArgumentException
 	{
+		long[] ogShape = tensor.dimensionsAsLongArray();
 		tensor = Utils.transpose(tensor);
 		PrimitiveBlocks< FloatType > blocks = PrimitiveBlocks.of( tensor );
 		long[] tensorShape = tensor.dimensionsAsLongArray();
@@ -231,8 +232,7 @@ private static Tensor<TFloat32> buildFloat(
 			sArr[i] = (int) tensorShape[i];
 		blocks.copy( new long[tensorShape.length], flatArr, sArr );
 		FloatDataBuffer dataBuffer = RawDataBufferFactory.create(flatArr, false);
-		Tensor<TFloat32> ndarray = TFloat32.tensorOf(Shape.of(tensor
-			.dimensionsAsLongArray()), dataBuffer);
+		Tensor<TFloat32> ndarray = TFloat32.tensorOf(Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}
 
@@ -250,6 +250,7 @@ private static Tensor<TFloat64> buildDouble(
 		RandomAccessibleInterval<DoubleType> tensor)
 		throws IllegalArgumentException
 	{
+		long[] ogShape = tensor.dimensionsAsLongArray();
 		tensor = Utils.transpose(tensor);
 		PrimitiveBlocks< DoubleType > blocks = PrimitiveBlocks.of( tensor );
 		long[] tensorShape = tensor.dimensionsAsLongArray();
@@ -261,8 +262,7 @@ private static Tensor<TFloat64> buildDouble(
 			sArr[i] = (int) tensorShape[i];
 		blocks.copy( new long[tensorShape.length], flatArr, sArr );
 		DoubleDataBuffer dataBuffer = RawDataBufferFactory.create(flatArr, false);
-		Tensor<TFloat64> ndarray = TFloat64.tensorOf(Shape.of(tensor
-			.dimensionsAsLongArray()), dataBuffer);
+		Tensor<TFloat64> ndarray = TFloat64.tensorOf(Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}
 }