diff --git a/python/solid_dmft/postprocessing/plot_correlated_bands.py b/python/solid_dmft/postprocessing/plot_correlated_bands.py
index 8cca5c75..42a2c842 100644
--- a/python/solid_dmft/postprocessing/plot_correlated_bands.py
+++ b/python/solid_dmft/postprocessing/plot_correlated_bands.py
@@ -245,8 +245,7 @@ def upscale(quantity, n_orb): return quantity * np.identity(n_orb)
         def invert_and_trace(w, eta, mu, e_mat, sigma, trace, proj=None):
             # inversion is automatically vectorized over first axis of 3D array (omega first index now)
             Glatt = np.linalg.inv(w + eta[None, ...] + mu[None, ...] - e_mat[None, ...] - sigma.transpose(2, 0, 1))
-            # TODO: fix this is only correct for diagonal Sigma
-            A_w_nu = -1.0/np.pi * np.diagonal(Glatt, axis1=1, axis2=2).imag
+            A_w_nu = -1.0/(2.0 * np.pi)* np.diagonal(Glatt - Glatt.transpose(0,2,1).conj(), axis1=1, axis2=2).imag
             if isinstance(proj, np.ndarray):
                 A_w_nu = A_w_nu * proj[None, :]
             if trace:
@@ -325,7 +324,7 @@ def upscale(quantity, n_orb): return quantity * np.identity(n_orb)
         def invert_and_trace(w, eta, mu, e_mat, sigma, proj=None):
             # inversion is automatically vectorized over first axis of 3D array (omega first index now)
             Glatt = np.linalg.inv(w + eta + mu - e_mat - sigma)
-            A_nu = -1.0/np.pi * np.diagonal(Glatt).imag
+            A_nu = -1.0/(2.0 * np.pi)* np.diagonal(Glatt - Glatt.transpose().conj()).imag
             if isinstance(proj, np.ndarray):
                 A_nu = A_nu * proj
             return np.sum(A_nu)