Skip to content

Commit 345dee2

Browse files
committed
simplify KMeans to use sklearn defaults
Signed-off-by: Weijia Chen <weijiac@nvidia.com>
1 parent 5aff0b3 commit 345dee2

1 file changed

Lines changed: 1 addition & 1 deletion

File tree

eval/video/build_benchmark_dataset.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ def _kmeans_select(
137137
uid_to_meta[uid] = json.load(fh)
138138

139139
logger.info(f"K-means K={num_clusters} ...")
140-
kmeans = KMeans(n_clusters=num_clusters, random_state=seed, n_init=10, max_iter=300)
140+
kmeans = KMeans(n_clusters=num_clusters, random_state=seed)
141141
labels = kmeans.fit_predict(emb_matrix)
142142
centers = kmeans.cluster_centers_
143143
cluster_sizes = np.bincount(labels, minlength=num_clusters)

0 commit comments

Comments
 (0)