diff --git a/coremltools/optimize/torch/_utils/k_means.py b/coremltools/optimize/torch/_utils/k_means.py
index 2e39810de..96213523a 100644
--- a/coremltools/optimize/torch/_utils/k_means.py
+++ b/coremltools/optimize/torch/_utils/k_means.py
@@ -675,10 +675,6 @@ def _cluster_weights_2d(
             max_iter=300,
         ).fit(weight_2d, sample_weight=importance_2d)
 
-        weight_2d.cpu()
-        if importance_2d is not None:
-            importance_2d.cpu()
-
         return kmeans_results.cluster_centers_.cpu(), kmeans_results.labels_.cpu()
 
     @classmethod