Use weights_only=True when loading sensitivities (#16)

BiswajeetRay7 · BiswajeetRay7 · commit 3520a7706460 · 2026-06-30T16:21:54.000-04:00
Signed-off-by: Biswajeet Ray &lt;raybiswajeet2@gmail.com&gt;
diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py
@@ -204,7 +204,7 @@ def prepare(
                 f"Loading precomputed sensitivities from {sensitivity_path} "
                 "for weighted k-means clustering"
             )
-            sensitivities = torch.load(sensitivity_path)
+            sensitivities = torch.load(sensitivity_path, weights_only=True)
             self._set_sensitivities_in_fake_palettize_modules(sensitivities)
 
         self._model.apply(_enable_observer)

Original file line number	Diff line number	Diff line change
`@@ -204,7 +204,7 @@ def prepare(`
`204`	`204`	`f"Loading precomputed sensitivities from {sensitivity_path} "`
`205`	`205`	`"for weighted k-means clustering"`
`206`	`206`	`)`
`207`		`- sensitivities = torch.load(sensitivity_path)`
	`207`	`+ sensitivities = torch.load(sensitivity_path, weights_only=True)`
`208`	`208`	`self._set_sensitivities_in_fake_palettize_modules(sensitivities)`
`209`	`209`
`210`	`210`	`self._model.apply(_enable_observer)`