diff --git a/pyscf/tdscf/_lr_eig.py b/pyscf/tdscf/_lr_eig.py
index 1dc58e649..2acd5e230 100644
--- a/pyscf/tdscf/_lr_eig.py
+++ b/pyscf/tdscf/_lr_eig.py
@@ -198,7 +198,7 @@ def eigh(aop, x0, precond, tol_residual=1e-5, lindep=1e-12, nroots=1,
         # remove subspace linear dependency
         for k, xk in enumerate(xt):
             if dx_norm[k] > tol_residual:
-                xt[k] = precond(xk, e[0])
+                xt[k] = precond(xk, w[k])
         xt -= xs.conj().dot(xt.T).T.dot(xs)
         xt_norm = np.linalg.norm(xt, axis=1)
 
@@ -426,7 +426,7 @@ def eig(aop, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=None,
         # remove subspace linear dependency
         for k, xk in enumerate(xt):
             if dx_norm[k] > tol_residual:
-                xt[k] = precond(xk, e[0])
+                xt[k] = precond(xk, w[k])
 
         xt -= xs.conj().dot(xt.T).T.dot(xs)
         c = _conj_dot(xs, xt)
diff --git a/pyscf/tdscf/rhf.py b/pyscf/tdscf/rhf.py
index 90b6260dc..2d33dba79 100644
--- a/pyscf/tdscf/rhf.py
+++ b/pyscf/tdscf/rhf.py
@@ -875,8 +875,14 @@ class TDBase(lib.StreamObject):
     def get_precond(self, hdiag):
         def precond(x, e, *args):
             if isinstance(e, numpy.ndarray):
-                e = e[0]
-            diagd = hdiag - (e-self.level_shift)
+                e = numpy.asarray(e)
+                if e.ndim == 0 or numpy.ndim(x) == 1:
+                    e = e.ravel()[0]
+                    diagd = hdiag - (e-self.level_shift)
+                else:
+                    diagd = hdiag - (e[:,None]-self.level_shift)
+            else:
+                diagd = hdiag - (e-self.level_shift)
             diagd[abs(diagd)<1e-8] = 1e-8
             return x/diagd
         return precond