diff --git a/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/shm/ShmBuilder.java b/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/shm/ShmBuilder.java
index ea45c1a..aff776d 100644
--- a/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/shm/ShmBuilder.java
+++ b/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/shm/ShmBuilder.java
@@ -111,11 +111,14 @@ private static void buildFromTensorFloat(Tensor tensor, String memoryName) throw
 			throw new IllegalArgumentException("Model output tensor with shape " + Arrays.toString(arrayShape) 
 					+ " is too big. Max number of elements per float output tensor supported: " + Integer.MAX_VALUE / 4);
 		SharedMemoryArray shma = SharedMemoryArray.readOrCreate(memoryName, arrayShape, new FloatType(), false, true);
-		long flatSize = 1;
+		/*long flatSize = 1;
     	for (long l : arrayShape) {flatSize *= l;}
     	ByteBuffer byteBuffer = ByteBuffer.allocate((int) (flatSize * Float.BYTES));
     	tensor.data_ptr_float().get(byteBuffer.asFloatBuffer().array());
         shma.getDataBufferNoHeader().put(byteBuffer);
+        */
+		RandomAccessibleInterval<?> rai = shma.getSharedRAI();
+		rai = ImgLib2Builder.build(tensor);
         if (PlatformDetection.isWindows()) shma.close();
     }