(fix): compute in iteratio

laminlabs · Sep 27, 2024 · 32fc41b · 32fc41b
1 parent ae0e5ff
commit 32fc41b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/figure_2_iteration_benchmark.py b/figure_2_iteration_benchmark.py
@@ -55,7 +55,7 @@ def _iterate(dataset, h5labels, random: bool = False, need_sort: bool = False):
     for batch_idx in index_iter(dataset.shape[0], BATCH_SIZE, shuffle=random):
         if random and need_sort:
             batch_idx.sort()
-        batch_X = dataset[batch_idx, :]
+        batch_X = dataset[batch_idx, :].compute()
         batch_labels = h5labels[batch_idx]