Remove unused A and B computation

antoine311200 · Mar 13, 2024 · 779d7d0 · 779d7d0
1 parent a6bc165
commit 779d7d0
Showing 1 changed file with 10 additions and 15 deletions.
diff --git a/galore_torch/galore_projector.py b/galore_torch/galore_projector.py
@@ -8,9 +8,9 @@ def __init__(self, rank, verbose=False, update_proj_gap=200, scale=1.0, proj_typ
         self.scale = scale
         self.ortho_matrix = None
         self.proj_type = proj_type
-        
+
     def project(self, full_rank_grad, iter):
-        
+
         if self.proj_type == 'std':
             if full_rank_grad.shape[0] >= full_rank_grad.shape[1]:
                 if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
@@ -41,7 +41,7 @@ def project(self, full_rank_grad, iter):
             if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
                 self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='full')
             low_rank_grad = torch.matmul(self.ortho_matrix[0].t(), full_rank_grad) @ self.ortho_matrix[1].t()
-                
+
         return low_rank_grad
 
     def project_back(self, low_rank_grad):
@@ -62,11 +62,11 @@ def project_back(self, low_rank_grad):
             full_rank_grad = torch.matmul(self.ortho_matrix, low_rank_grad)
         elif self.proj_type == 'full':
             full_rank_grad = torch.matmul(self.ortho_matrix[0], low_rank_grad) @ self.ortho_matrix[1]
-        
-        
+
+
         return full_rank_grad * self.scale
-        
-        
+
+
     # svd decomposition
     def get_orthogonal_matrix(self, weights, rank, type):
         module_params = weights
@@ -79,20 +79,18 @@ def get_orthogonal_matrix(self, weights, rank, type):
         else:
             float_data = True
             matrix = module_params.data
-            
+
         U, s, Vh = torch.linalg.svd(matrix, full_matrices = False)
-        
+
         #make the smaller matrix always to be orthogonal matrix
         if type=='right':
-            A = U[:, :rank] @ torch.diag(s[:rank])
             B = Vh[:rank, :]
-            
+
             if not float_data:
                 B = B.to(original_device).type(original_type)
             return B
         elif type=='left':
             A = U[:, :rank]
-            B = torch.diag(s[:rank]) @ Vh[:rank, :]
             if not float_data:
                 A = A.to(original_device).type(original_type)
             return A
@@ -105,6 +103,3 @@ def get_orthogonal_matrix(self, weights, rank, type):
             return [A, B]
         else:
             raise ValueError('type should be left, right or full')
-
-
-