fix register_buffer in MimiEuclideanCodebook (#35759)

Co-authored-by: eustlb <94853470+eustlb@users.noreply.github.com>
huggingface · Jan 20, 2025 · 56afd2f · 56afd2f
1 parent abe57b6
commit 56afd2f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/models/mimi/modeling_mimi.py b/src/transformers/models/mimi/modeling_mimi.py
@@ -1254,7 +1254,7 @@ def __init__(self, config: MimiConfig, epsilon: float = 1e-5):
 
         self.codebook_size = config.codebook_size
 
-        self.register_buffer("initialized", torch.Tensor([True]))
+        self.register_buffer("initialized", torch.tensor([True], dtype=torch.float32))
         self.register_buffer("cluster_usage", torch.ones(config.codebook_size))
         self.register_buffer("embed_sum", embed)
         self._embed = None