From f9e7a41915bce8fe02fb29b87565b08d28424c5e Mon Sep 17 00:00:00 2001
From: Weston Pace <weston.pace@gmail.com>
Date: Tue, 2 Jul 2024 17:03:56 -0700
Subject: [PATCH] Fix up call to compute_partitions in create_index

---
 python/python/lance/dataset.py | 2 +-
 python/python/lance/indices.py | 2 +-
 python/python/lance/vector.py  | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py
index a7efb7c794..7bb33c2eaf 100644
--- a/python/python/lance/dataset.py
+++ b/python/python/lance/dataset.py
@@ -1490,7 +1490,7 @@ def create_index(
                     accelerator,
                 )
                 partitions_file = compute_partitions(
-                    self, column, kmeans, batch_size=20480
+                    self, column[0], kmeans, batch_size=20480
                 )
                 kwargs["precomputed_partitions_file"] = partitions_file
 
diff --git a/python/python/lance/indices.py b/python/python/lance/indices.py
index 81a71f6d2d..8b5178bb1f 100644
--- a/python/python/lance/indices.py
+++ b/python/python/lance/indices.py
@@ -168,7 +168,7 @@ def train_ivf(
             # Use accelerator to train ivf centroids
             from .vector import train_ivf_centroids_on_accelerator
 
-            ivf_centroids, kmeans = train_ivf_centroids_on_accelerator(
+            ivf_centroids, _ = train_ivf_centroids_on_accelerator(
                 self.dataset,
                 column[0],
                 num_partitions,
diff --git a/python/python/lance/vector.py b/python/python/lance/vector.py
index 4d994fcc95..dc7de7d09d 100644
--- a/python/python/lance/vector.py
+++ b/python/python/lance/vector.py
@@ -183,7 +183,7 @@ def train_ivf_centroids_on_accelerator(
         np.save(f, centroids)
     logging.info("Saved centroids to %s", f.name)
 
-    return centroids, compute_partitions(dataset, column, kmeans, batch_size=20480)
+    return centroids, kmeans
 
 
 def compute_partitions(