Updating depricated torch.trtrs (#54)

VincentStimper · Dec 4, 2023 · 27d4bf8 · 27d4bf8
1 parent 8272cdc
commit 27d4bf8
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/normflows/flows/mixing.py b/normflows/flows/mixing.py
@@ -505,9 +505,9 @@ def weight_inverse(self):
         """
         lower, upper = self._create_lower_upper()
         identity = torch.eye(self.features, self.features)
-        lower_inverse, _ = torch.trtrs(identity, lower, upper=False, unitriangular=True)
-        weight_inverse, _ = torch.trtrs(
-            lower_inverse, upper, upper=True, unitriangular=False
+        lower_inverse = torch.linalg.solve_triangular(lower, identity, upper=False, unitriangular=True)
+        weight_inverse = torch.linalg.solve_triangular(
+            upper, lower_inverse, upper=True, unitriangular=False
         )
         return weight_inverse