diff --git a/arbiter/arbiter-ui/src/test/java/org/deeplearning4j/arbiter/optimize/TestBasic.java b/arbiter/arbiter-ui/src/test/java/org/deeplearning4j/arbiter/optimize/TestBasic.java
index 08f130369..3ecefe0b3 100644
--- a/arbiter/arbiter-ui/src/test/java/org/deeplearning4j/arbiter/optimize/TestBasic.java
+++ b/arbiter/arbiter-ui/src/test/java/org/deeplearning4j/arbiter/optimize/TestBasic.java
@@ -57,12 +57,12 @@ import org.deeplearning4j.ui.api.UIServer;
 import org.deeplearning4j.ui.model.storage.InMemoryStatsStorage;
 import org.junit.Ignore;
 import org.junit.Test;
+import org.nd4j.common.function.Function;
 import org.nd4j.evaluation.classification.Evaluation;
 import org.nd4j.linalg.activations.Activation;
 import org.nd4j.linalg.api.buffer.DataType;
 import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
 import org.nd4j.linalg.factory.Nd4j;
-import org.nd4j.linalg.function.Function;
 import org.nd4j.linalg.lossfunctions.LossFunctions;
 
 import java.io.File;