Only cast if necessary

VincentStimper · Oct 26, 2023 · 912891d · 912891d
1 parent a2b5805
commit 912891d
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/normflows/utils/splines.py b/normflows/utils/splines.py
@@ -87,8 +87,12 @@ def unconstrained_rational_quadratic_spline(
         min_bin_height=min_bin_height,
         min_derivative=min_derivative,
     )
-    outputs[inside_interval_mask] = outputs_masked.to(outputs.dtype)
-    logabsdet[inside_interval_mask] = logabsdet_masked.to(logabsdet.dtype)
+    if outputs.dtype == outputs_masked.dtype and logabsdet.dtype == logabsdet_masked.dtype:
+        outputs[inside_interval_mask] = outputs_masked
+        logabsdet[inside_interval_mask] = logabsdet_masked
+    else:
+        outputs[inside_interval_mask] = outputs_masked.to(outputs.dtype)
+        logabsdet[inside_interval_mask] = logabsdet_masked.to(logabsdet.dtype)
 
     return outputs, logabsdet