Don't require the user to set the sign of h_init for DPM-Solver-Adaptive

crowsonkb · Sep 30, 2022 · 99f1035 · 99f1035
1 parent d76c63a
commit 99f1035
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/k_diffusion/sampling.py b/k_diffusion/sampling.py
@@ -352,11 +352,7 @@ def dpm_solver_adaptive(self, x, t_start, t_end, order=3, rtol=0.05, atol=0.0078
         if order not in {2, 3}:
             raise ValueError('order should be 2 or 3')
         forward = t_end > t_start
-        if forward and h_init <= 0:
-            raise ValueError('For forward ODE integration, h_init must be positive')
-        if not forward and h_init >= 0:
-            raise ValueError('For reverse ODE integration, h_init must be negative')
-
+        h_init = abs(h_init) * (1 if forward else -1)
         atol = torch.tensor(atol)
         rtol = torch.tensor(rtol)
         s = t_start