Update modeling_wavcoch.py

modeling_wavcoch.py

@@ -508,7 +508,7 @@ class WavCoch(PreTrainedModel):
             if pad:
                 wav = F.pad(wav, (self.N - self.hop_length, 0), mode="constant", value=0.0)
 
-            codes = self.quantize(wav, pad=False).long()  # already padded
+            codes = self.quantize(wav, pad=False)  # already padded
             return BatchEncoding({"input_values": codes, "input_ids": codes})
 
         # Training / reconstruction mode
@@ -542,7 +542,7 @@ class WavCoch(PreTrainedModel):
         x = real_part + imag_part
         x = self.encoder(x).permute(0, 2, 1)        # (B, T, D)
         _, indices = self.quantizer(x)              # (B, T)
-        return indices
+        return indices.long()
 
     @torch.no_grad()
     def decode(self, indices: torch.Tensor) -> torch.Tensor: