diff --git a/src/main/java/tensorflow/v1/Tensorflow1Interface.java b/src/main/java/tensorflow/v1/Tensorflow1Interface.java
index a0af331fd58be6946174c86e1ef4e37c37825900..bb68bb70adb83bbd56379f1965fbe4cd11d4757a 100644
--- a/src/main/java/tensorflow/v1/Tensorflow1Interface.java
+++ b/src/main/java/tensorflow/v1/Tensorflow1Interface.java
@@ -7,15 +7,16 @@ import org.bioimageanalysis.icy.deeplearning.exceptions.LoadModelException;
 import org.bioimageanalysis.icy.deeplearning.exceptions.RunModelException;
 import org.bioimageanalysis.icy.deeplearning.tensor.Tensor;
 import org.bioimageanalysis.icy.deeplearning.utils.DeepLearningInterface;
-import org.bioimageanalysis.icy.tensorflow.v1.tensor.Nd4fBuilder;
+import org.bioimageanalysis.icy.tensorflow.v1.tensor.ImgLib2Builder;
 import org.bioimageanalysis.icy.tensorflow.v1.tensor.TensorBuilder;
-import org.nd4j.shade.protobuf.InvalidProtocolBufferException;
 import org.tensorflow.SavedModelBundle;
 import org.tensorflow.Session;
 import org.tensorflow.framework.MetaGraphDef;
 import org.tensorflow.framework.SignatureDef;
 import org.tensorflow.framework.TensorInfo;
 
+import com.google.protobuf.InvalidProtocolBufferException;
+
 
 /**
  * This plugin includes the libraries to convert back and forth TensorFlow 1 to Sequences and IcyBufferedImages.
@@ -120,7 +121,7 @@ public class Tensorflow1Interface implements DeepLearningInterface
 		if (outputNDArrays.size() != outputTensors.size())
 			throw new RunModelException(outputNDArrays.size(), outputTensors.size());
 		for (int i = 0; i < outputNDArrays.size(); i ++) {
-			outputTensors.get(i).setData(Nd4fBuilder.build(outputNDArrays.get(i)));
+			outputTensors.get(i).setData(ImgLib2Builder.build(outputNDArrays.get(i)));
 		}
 		return outputTensors;
 	}