Update stabilize_derivative

StatMixedML · Aug 28, 2023 · b3e83a9 · b3e83a9
1 parent 923e11e
commit b3e83a9
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/lightgbmlss/distributions/mixture_distribution_utils.py b/lightgbmlss/distributions/mixture_distribution_utils.py
@@ -596,6 +596,13 @@ def stabilize_derivative(self, input_der: torch.Tensor, type: str = "MAD") -> to
             div = torch.where(div > torch.tensor(10000.0), torch.tensor(10000.0), div)
             stab_der = input_der / div
 
+        if type == "None":
+            stab_der = torch.nan_to_num(input_der,
+                                        nan=float(torch.nanmean(input_der)),
+                                        posinf=float(torch.nanmean(input_der)),
+                                        neginf=float(torch.nanmean(input_der))
+                                        )
+
         return stab_der
 
     def dist_select(self,