diff --git a/scripts/sortcolor.py b/scripts/sortcolor.py
index 8ceb975..d7abfa9 100644
--- a/scripts/sortcolor.py
+++ b/scripts/sortcolor.py
@@ -35,7 +35,13 @@ FONTSIZE = args.fontsize
 INNER_RADIUS = args.radius
 DIR = "/teamspace/studios/this_studio/out_sortcolors"
 
-Path(DIR).mkdir(exist_ok=True, parents=True)
+
+prefix = ""
+if KIND == "umap":
+    prefix = f"{SEED:04d}"
+FDIR = f"{DIR}/{KIND}"
+Path(FDIR).mkdir(exist_ok=True, parents=True)
+fname = f"{FDIR}/{prefix}_sorted_colors_circle.png"
 
 
 def peano_curve(n):
@@ -133,9 +139,10 @@ sorted_indices = np.argsort(preds)
 
 # Save the sorted indices to disk
 # if (KIND == "umap") or (KIND != "umap"):
-Path(f"scripts/{KIND}").mkdir(parents=True, exist_ok=True)
-file_path = f"scripts/{KIND}/preds_{SEED:06d}.npy"
-np.save(file_path, preds)
+PDIR = f"scripts/{KIND}"
+Path(PDIR).mkdir(parents=True, exist_ok=True)
+file_path = f"{PDIR}/{SEED:06d}.npy"
+np.save(file_path, preds.ravel())
 print(f"Predictions saved to {file_path}")
 
 # Sort colors by the 1D representation
@@ -281,11 +288,6 @@ def plot_preds(
     plt.close()
 
 
-prefix = ""
-if KIND == "umap":
-    prefix = f"{SEED:04d}"
-fname = f"{DIR}/{prefix}{KIND}_sorted_colors_circle.png"
-
 plot_preds(
     preds,
     np.array(rgb_values),