Add hooks to ensure saving of kmeans model works correctly (#3047)

pplantinga · web-flow · commit 9c826ef3937b · 2026-04-03T16:45:42.000+01:00
diff --git a/speechbrain/integrations/audio_tokenizers/kmeans.py b/speechbrain/integrations/audio_tokenizers/kmeans.py
@@ -7,7 +7,10 @@
 import joblib
 import torch
 
+from speechbrain.utils import checkpoints
 
+
+@checkpoints.register_checkpoint_hooks
 class MiniBatchKMeansSklearn(torch.nn.Module):
     """A wrapper for scikit-learn MiniBatchKMeans, providing integration with PyTorch tensors.
 
@@ -64,6 +67,7 @@ def to(self, device=None, **kwargs):
         self.device = device
         return super().to(device)
 
+    @checkpoints.mark_as_saver
     def save(self, path):
         """Saves the model to the specified file.
 
@@ -74,6 +78,7 @@ def save(self, path):
         """
         joblib.dump(self.kmeans, path)
 
+    @checkpoints.mark_as_loader
     def load(self, path, end_of_epoch):
         """Loads the model from the specified file.
 
