diff --git a/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/TsneTest.java b/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/TsneTest.java
index 7edb8ea3e..20ba99cbe 100644
--- a/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/TsneTest.java
+++ b/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/TsneTest.java
@@ -30,6 +30,7 @@ import org.deeplearning4j.nn.conf.WorkspaceMode;
 
 import org.junit.jupiter.api.Test;
 
+import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.api.io.TempDir;
 import org.nd4j.linalg.api.buffer.DataType;
 import org.nd4j.linalg.api.ndarray.INDArray;
@@ -42,13 +43,9 @@ import java.util.ArrayList;
 import java.util.List;
 
 @Slf4j
+@Timeout(300)
 public class TsneTest extends BaseDL4JTest {
 
-    @Override
-    public long getTimeoutMilliseconds() {
-        return 180000L;
-    }
-
     @TempDir
     public File testDir;
 
diff --git a/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/models/paragraphvectors/ParagraphVectorsTest.java b/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/models/paragraphvectors/ParagraphVectorsTest.java
index cac34901f..84869467b 100644
--- a/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/models/paragraphvectors/ParagraphVectorsTest.java
+++ b/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/models/paragraphvectors/ParagraphVectorsTest.java
@@ -78,11 +78,6 @@ import static org.junit.jupiter.api.Assertions.*;
 @Timeout(240)
 public class ParagraphVectorsTest extends BaseDL4JTest {
 
-    @Override
-    public long getTimeoutMilliseconds() {
-        return isIntegrationTests() ? 600_000 : 240_000;
-    }
-
     @TempDir
     public File testDir;
 
diff --git a/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/models/word2vec/Word2VecTestsSmall.java b/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/models/word2vec/Word2VecTestsSmall.java
index 6d7bfaf63..19681185c 100644
--- a/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/models/word2vec/Word2VecTestsSmall.java
+++ b/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/models/word2vec/Word2VecTestsSmall.java
@@ -60,14 +60,10 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
 
 
 @Slf4j
+@Timeout(300)
 public class Word2VecTestsSmall extends BaseDL4JTest {
     WordVectors word2vec;
 
-    @Override
-    public long getTimeoutMilliseconds() {
-        return isIntegrationTests() ? 240000 : 60000;
-    }
-
     @BeforeEach
     public void setUp() throws Exception {
         word2vec = WordVectorSerializer.readWord2VecModel(new ClassPathResource("vec.bin").getFile());
diff --git a/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/models/word2vec/iterator/Word2VecDataSetIteratorTest.java b/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/models/word2vec/iterator/Word2VecDataSetIteratorTest.java
index a72f92211..7d806aafb 100644
--- a/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/models/word2vec/iterator/Word2VecDataSetIteratorTest.java
+++ b/cavis-dnn/cavis-dnn-nlp/src/test/java/org/deeplearning4j/models/word2vec/iterator/Word2VecDataSetIteratorTest.java
@@ -34,6 +34,7 @@ import org.deeplearning4j.text.tokenization.tokenizerfactory.DefaultTokenizerFac
 import org.deeplearning4j.text.tokenization.tokenizerfactory.TokenizerFactory;
 
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
 import org.nd4j.linalg.api.ndarray.INDArray;
 import org.nd4j.linalg.dataset.DataSet;
 import org.nd4j.common.resources.Resources;
@@ -46,13 +47,9 @@ import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 
+@Timeout(300)
 public class Word2VecDataSetIteratorTest extends BaseDL4JTest {
 
-    @Override
-    public long getTimeoutMilliseconds() {
-        return 60000L;
-    }
-
     /**
      * Basically all we want from this test - being able to finish without exceptions.
      */