update gpu

app.py

@@ -555,7 +555,7 @@ def compute_ncut(
         affinity_focal_gamma=affinity_focal_gamma,
         knn=knn_ncut,
     ).fit_transform(features.reshape(-1, features.shape[-1]))
-    print(f"NCUT time (cpu): {time.time() - start:.2f}s")
+    print(f"NCUT time: {time.time() - start:.2f}s")
     
     start = time.time()
     if embedding_method == "UMAP":
@@ -564,7 +564,7 @@ def compute_ncut(
             n_neighbors=n_neighbors,
             min_dist=min_dist,
         )
-        print(f"UMAP time (cpu): {time.time() - start:.2f}s")
+        print(f"UMAP time: {time.time() - start:.2f}s")
     elif embedding_method == "t-SNE":    
         X_3d, rgb = rgb_from_tsne_3d(
             eigvecs,
@@ -572,7 +572,7 @@ def compute_ncut(
             perplexity=perplexity,
             knn=knn_tsne,
         )
-        print(f"t-SNE time (cpu): {time.time() - start:.2f}s")
+        print(f"t-SNE time: {time.time() - start:.2f}s")
     else:
         raise ValueError(f"Embedding method {embedding_method} not supported.")
     
@@ -657,7 +657,7 @@ demo = gr.Interface(
     additional_inputs=[
         gr.Dropdown(["attn", "mlp", "block"], label="Node type", value="block", elem_id="node_type", info="attn: attention output, mlp: mlp output, block: sum of residual stream"),
         gr.Slider(0.01, 1, step=0.01, label="Affinity focal gamma", value=0.3, elem_id="affinity_focal_gamma", info="decrease for more aggressive cleaning on the affinity matrix"),
-        gr.Slider(100, 10000, step=100, label="num_sample (NCUT)", value=5000, elem_id="num_sample_ncut", info="for Nyström approximation"),
+        gr.Slider(100, 50000, step=100, label="num_sample (NCUT)", value=25000, elem_id="num_sample_ncut", info="for Nyström approximation"),
         gr.Slider(1, 100, step=1, label="KNN (NCUT)", value=10, elem_id="knn_ncut", info="for Nyström approximation"),
         gr.Dropdown(["t-SNE", "UMAP"], label="Embedding method", value="t-SNE", elem_id="embedding_method"),
         gr.Slider(100, 1000, step=100, label="num_sample (t-SNE/UMAP)", value=300, elem_id="num_sample_tsne", info="for Nyström approximation. Adding will slow down quite a lot"),