implement torch.frac() on-GPU.

pytorch/pytorch#77764 (comment) pytorch/pytorch#77764 (comment)
Birch-san · Nov 5, 2022 · c56a015 · c56a015
1 parent 0602097
commit c56a015
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/k_diffusion/external.py b/k_diffusion/external.py
@@ -79,7 +79,9 @@ def sigma_to_t(self, sigma, quantize=None):
 
     def t_to_sigma(self, t):
         t = t.float()
-        low_idx, high_idx, w = t.floor().long(), t.ceil().long(), t.frac()
+        low_idx = t.floor().long()
+        high_idx = t.ceil().long()
+        w = t-low_idx if t.device.type == 'mps' else t.frac()
         log_sigma = (1 - w) * self.log_sigmas[low_idx] + w * self.log_sigmas[high_idx]
         return log_sigma.exp()