Fix version error

dataflowr · Jan 10, 2023 · 9bd7d2d · 9bd7d2d
1 parent 29f0b96
commit 9bd7d2d
Showing 1 changed file with 14 additions and 6 deletions.
diff --git a/normflows/flows/mixing.py b/normflows/flows/mixing.py
@@ -451,12 +451,20 @@ def inverse_no_cache(self, inputs):
         """
         lower, upper = self._create_lower_upper()
         outputs = inputs - self.bias
-        outputs = torch.linalg.solve_triangular(
-            lower, outputs.t(), upper=False, unitriangular=True
-        )
-        outputs = torch.linalg.solve_triangular(
-            upper, outputs, upper=True, unitriangular=False
-        )
+        try:
+            outputs = torch.linalg.solve_triangular(
+                lower, outputs.t(), upper=False, unitriangular=True
+            )
+            outputs = torch.linalg.solve_triangular(
+                upper, outputs, upper=True, unitriangular=False
+            )
+        except:
+            outputs, _ = torch.triangular_solve(
+                outputs.t(), lower, upper=False, unitriangular=True
+            )
+            outputs, _ = torch.triangular_solve(
+                outputs, upper, upper=True, unitriangular=False
+            )
         outputs = outputs.t()
 
         logabsdet = -self.logabsdet()