diff --git a/task1.py b/task1.py
index 415a3c3..ffb666b 100644
--- a/task1.py
+++ b/task1.py
@@ -286,7 +286,11 @@ def task1(
 
     k = 30
 
-    for nprobe in range(1, 30 + 1):
+    nprobes = range(1, 30 + 1)
+    if dataset_size == '300K':
+        nprobes = [1]
+
+    for nprobe in nprobes:
         start = time.time()
         D, I = lmi.search(queries, k, nprobe)
         searchtime = time.time() - start
diff --git a/task2.py b/task2.py
index 599330a..f80dc44 100644
--- a/task2.py
+++ b/task2.py
@@ -347,7 +347,11 @@ def task2(
 
     k = 30
 
-    for nprobe in [1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 80, 90, 100]:
+    nprobes = [1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 80, 90, 100]
+    if dataset_size == '300K':
+        nprobes = [1]
+
+    for nprobe in nprobes:
         start = time.time()
         D, I = lmi.search(queries, decomposed_queries, k, nprobe, ncandidates)
         searchtime = time.time() - start
diff --git a/task3.py b/task3.py
index c698c0d..ddb4487 100644
--- a/task3.py
+++ b/task3.py
@@ -323,7 +323,11 @@ def task3(
 
     k = 30
 
-    for nprobe in range(1, 30 + 1):
+    nprobes = range(1, 30 + 1)
+    if dataset_size == '300K':
+        nprobes = [1]
+
+    for nprobe in nprobes:
         start = time.time()
         D, I = lmi.search(queries, decomposed_queries, k, nprobe)
         searchtime = time.time() - start