Increase float precision in residual FSQ

Author: julian-parker

stable_codec/residual_fsq.py

@@ -18,14 +18,15 @@ def __init__(self, stages: List[Tuple[List[int], float]]):
         self.codebook_size = sum(map(len, stages)) * self.n_codebooks
 
     def encode(self, x):
-        z = torch.tanh(x)
+        input_dtype = x.dtype
+        z = torch.tanh(x.to(torch.float64))
         z = rearrange(z, "b c n -> b n c")
 
         r = z
         res_ids = []
         for quantizer in self.quantizers:
             q, ids = quantizer(r, skip_tanh=True)
-            r = r - q
+            r = r - q.to(torch.float64)
             res_ids.append(ids)
 
         return res_ids