Revert "Unable to use MimiModel with DeepSpeed ZeRO-3" (huggingface…

…#35755) Revert "Unable to use `MimiModel` with DeepSpeed ZeRO-3 (huggingface#34735)" This reverts commit 54fd7e9.
neuralmagic · Jan 17, 2025 · 936a731 · 936a731
1 parent 10e8cd0
commit 936a731
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]))
         self.register_buffer("cluster_usage", torch.ones(config.codebook_size))
         self.register_buffer("embed_sum", embed)
         self._embed = None