fix validate noise schedule for training

Author: arrjon

bayesflow/experimental/diffusion_model.py

@@ -58,7 +58,7 @@ def scale_base_distribution(self):
             raise ValueError(f"Unknown variance type: {self.variance_type}")
 
     @abstractmethod
-    def get_log_snr(self, t: Tensor, training: bool) -> Tensor:
+    def get_log_snr(self, t: Union[float, Tensor], training: bool) -> Tensor:
         """Get the log signal-to-noise ratio (lambda) for a given diffusion time."""
         pass
 
@@ -145,10 +145,10 @@ def validate(self):
                 raise ValueError("t(0) must be finite.")
             if not ops.isfinite(self.get_t_from_log_snr(self._log_snr_min, training=training)):
                 raise ValueError("t(1) must be finite.")
-            if not ops.isfinite(self.derivative_log_snr(self._log_snr_max, training=training)):
-                raise ValueError("dt/t log_snr(0) must be finite.")
-            if not ops.isfinite(self.derivative_log_snr(self._log_snr_min, training=training)):
-                raise ValueError("dt/t log_snr(1) must be finite.")
+        if not ops.isfinite(self.derivative_log_snr(self._log_snr_max, training=False)):
+            raise ValueError("dt/t log_snr(0) must be finite.")
+        if not ops.isfinite(self.derivative_log_snr(self._log_snr_min, training=False)):
+            raise ValueError("dt/t log_snr(1) must be finite.")
 
 
 @serializable
@@ -168,7 +168,7 @@ def __init__(self, min_log_snr: float = -15, max_log_snr: float = 15):
         self._t_min = self.get_t_from_log_snr(log_snr_t=self._log_snr_max, training=True)
         self._t_max = self.get_t_from_log_snr(log_snr_t=self._log_snr_min, training=True)
 
-    def get_log_snr(self, t: Tensor, training: bool) -> Tensor:
+    def get_log_snr(self, t: Union[float, Tensor], training: bool) -> Tensor:
         """Get the log signal-to-noise ratio (lambda) for a given diffusion time."""
         t_trunc = self._t_min + (self._t_max - self._t_min) * t
         # SNR = -log(exp(t^2) - 1)
@@ -228,7 +228,7 @@ def __init__(self, min_log_snr: float = -15, max_log_snr: float = 15, s_shift_co
         self._t_min = self.get_t_from_log_snr(log_snr_t=self._log_snr_max, training=True)
         self._t_max = self.get_t_from_log_snr(log_snr_t=self._log_snr_min, training=True)
 
-    def get_log_snr(self, t: Tensor, training: bool) -> Tensor:
+    def get_log_snr(self, t: Union[float, Tensor], training: bool) -> Tensor:
         """Get the log signal-to-noise ratio (lambda) for a given diffusion time."""
         t_trunc = self._t_min + (self._t_max - self._t_min) * t
         # SNR = -2 * log(tan(pi*t/2))
@@ -289,7 +289,7 @@ def __init__(self, sigma_data: float = 0.5, sigma_min: float = 0.002, sigma_max:
         self._t_min = self.get_t_from_log_snr(log_snr_t=self._log_snr_max, training=True)
         self._t_max = self.get_t_from_log_snr(log_snr_t=self._log_snr_min, training=True)
 
-    def get_log_snr(self, t: Tensor, training: bool) -> Tensor:
+    def get_log_snr(self, t: Union[float, Tensor], training: bool) -> Tensor:
         """Get the log signal-to-noise ratio (lambda) for a given diffusion time."""
         t_trunc = self._t_min + (self._t_max - self._t_min) * t
         if training: