From 1b450ee15d71d694299a6acc06f6372a8ad37983 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Patrick=20Gel=C3=9F?=
 <38036185+PGelss@users.noreply.github.com>
Date: Sun, 2 Jun 2024 19:32:59 +0200
Subject: [PATCH] fix in diag

---
 scikit_tt/tensor_train.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/scikit_tt/tensor_train.py b/scikit_tt/tensor_train.py
index 88d75be..45e0fb0 100644
--- a/scikit_tt/tensor_train.py
+++ b/scikit_tt/tensor_train.py
@@ -1704,14 +1704,12 @@ def pinv(self,
 
         return p_inv
     
-    def diag(self, t: 'TT', diag_list: List) -> 'TT':
+    def diag(self, diag_list: List) -> 'TT':
         """
         Construction of diagonal MPO from MPS.
 
         Parameters
         ----------
-        t: TT
-            MPS decomposition
         diag_list: list
             cores to diagonalize
 
@@ -1721,7 +1719,7 @@ def diag(self, t: 'TT', diag_list: List) -> 'TT':
             MPO decomposition
         """
 
-        cores = t.cores.copy()
+        cores = self.cores.copy()
 
         for i in diag_list:
             r1, d, _, r2 = cores[i].shape