Pivoted cholesky checks that max_iter isn't more than matrix size
@@ -31,6 +31,9 @@ def pivoted_cholesky(matrix, max_iter, error_tol=1e-3):
matrix_diag.unsqueeze_(0)
# matrix_diag is now batch_size x n
+ # Make sure max_iter isn't bigger than the matrix
+ max_iter = min(max_iter, matrix_size)
+
errors = torch.norm(matrix_diag, 1, dim=1)
permutation = matrix_diag.new(matrix_size).long()
torch.arange(0, matrix_size, out=permutation)
@@ -70,6 +73,7 @@ def pivoted_cholesky(matrix, max_iter, error_tol=1e-3):
if isinstance(row, torch.autograd.Variable):
row = row.data
+ if m + 1 < matrix_size:
pi_i = permutation[:, m + 1:]
L_m_new = row.gather(1, pi_i)
if m > 0:
