diff --git a/src/ivfkmeans.c b/src/ivfkmeans.c
index 1c73225..31eb2a3 100644
--- a/src/ivfkmeans.c
+++ b/src/ivfkmeans.c
@@ -5,10 +5,7 @@
 
 #include "ivfflat.h"
 #include "miscadmin.h"
-
-#ifdef IVFFLAT_MEMORY
 #include "utils/memutils.h"
-#endif
 
 /*
  * Initialize with kmeans++
@@ -201,6 +198,8 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
 	float	   *s;
 	float	   *halfcdist;
 	float	   *newcdist;
+	MemoryContext kmeansCtx;
+	MemoryContext oldCtx;
 
 	/* Calculate allocation sizes */
 	Size		samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->dim);
@@ -234,6 +233,12 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
 	normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
 	collation = index->rd_indcollation[0];
 
+	/* Use memory context */
+	kmeansCtx = AllocSetContextCreate(CurrentMemoryContext,
+									  "Ivfflat kmeans temporary context",
+									  ALLOCSET_DEFAULT_SIZES);
+	oldCtx = MemoryContextSwitchTo(kmeansCtx);
+
 	/* Allocate space */
 	/* Use float instead of double to save memory */
 	centerCounts = palloc(centerCountsSize);
@@ -472,14 +477,8 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
 			break;
 	}
 
-	VectorArrayFree(newCenters);
-	pfree(centerCounts);
-	pfree(closestCenters);
-	pfree(lowerBound);
-	pfree(upperBound);
-	pfree(s);
-	pfree(halfcdist);
-	pfree(newcdist);
+	MemoryContextSwitchTo(oldCtx);
+	MemoryContextDelete(kmeansCtx);
 }
 
 /*