diff --git a/codes/hf.py b/codes/hf.py
index ec8a646b9b44e2489b47b0481796771232f4344e..e26fa58924eb5f52891e5c5980494489cb5f9cb3 100644
--- a/codes/hf.py
+++ b/codes/hf.py
@@ -163,7 +163,8 @@ def find_groundstate_ham(
     cost_function=default_cost,
     guess=None,
     optimizer=anderson,
-    optimizer_kwargs=None
+    optimizer_kwargs={},
+    return_model=False
 ):
     """
     Self-consistent loop to find groundstate Hamiltonian.
@@ -192,7 +193,7 @@ def find_groundstate_ham(
         dim=dim,
         return_ks=True
     )
-    H_int = utils.kgrid_hamiltonian(nk, Vk, dim=dim, return_ks=True)
+    H_int = utils.kgrid_hamiltonian(nk=nk, hk=Vk, dim=dim)
     vectors = utils.generate_vectors(cutoff, dim)
     if guess is None:
         guess = utils.generate_guess(
@@ -201,16 +202,18 @@ def find_groundstate_ham(
             scale=np.max(np.abs(H_int))
         )
     guess_k = utils.kgrid_hamiltonian(
-        nk,
-        utils.model2hk(guess),
-        return_ks=True
+        nk=nk,
+        hk=utils.model2hk(guess),
+        dim=dim,
     )
     cost_function_wrapped = partial(
         cost_function,
         hamiltonians_0=hamiltonians_0,
         H_int=H_int,
         filling=filling,
-        vectors=vectors
     )
-    mf_k = optimizer(cost_function_wrapped, guess_k, **optimizer_kwargs)
-    return utils.hk2tb_model(mf_k, vectors, ks)
+    if return_model:
+        mf_k = optimizer(cost_function_wrapped, guess_k, **optimizer_kwargs)
+        return mf_k, utils.hk2tb_model(hamiltonians_0 + mf_k, vectors, ks)
+    else:
+        return optimizer(cost_function_wrapped, guess_k, **optimizer_kwargs)
diff --git a/codes/utils.py b/codes/utils.py
index 6d761b1874a467303d541dc31f6a001548c4086c..95b68ff6d9c0dde6d95381e3213fc636eb0cb0d6 100644
--- a/codes/utils.py
+++ b/codes/utils.py
@@ -171,16 +171,16 @@ def model2hk(tb_model):
             ham += tb_model[vector] * np.exp(
                 1j * np.dot(k, np.array(vector, dtype=float))
             )
-            if np.linalg.norm(np.array(vector)) > 0:
-                ham += tb_model[vector].T.conj() * np.exp(
-                    -1j * np.dot(k, np.array(vector))
-                )
+            # if np.linalg.norm(np.array(vector)) > 0:
+            #     ham += tb_model[vector].T.conj() * np.exp(
+            #         -1j * np.dot(k, np.array(vector))
+            #     )
         return ham
 
     return bloch_ham
 
 
-def kgrid_hamiltonian(nk, hk, dim, return_info=False):
+def kgrid_hamiltonian(nk, hk, dim, return_ks=False):
     """
     Evaluates Hamiltonian on a k-point grid.
 
@@ -190,8 +190,8 @@ def kgrid_hamiltonian(nk, hk, dim, return_info=False):
         Number of k-points along each direction.
     hk : function
         Calculates the Hamiltonian at a given k-point.
-    return_info : bool
-        If `True`, returns k-points and dimension of the tight.
+    return_ks : bool
+        If `True`, returns k-points.
 
     Returns:
     --------
@@ -213,7 +213,7 @@ def kgrid_hamiltonian(nk, hk, dim, return_info=False):
         ham.append(hk(k))
     ham = np.array(ham)
     shape = (*[nk] * dim, ham.shape[-1], ham.shape[-1])
-    if return_info:
+    if return_ks:
         return ham.reshape(*shape), ks
     else:
         return ham.reshape(*shape)
@@ -272,7 +272,7 @@ def generate_guess(vectors, ndof, scale=0.1):
     return guess
 
 def generate_vectors(cutoff, dim):
-    return np.array([*product(*([[*range(-cutoff, cutoff)]]*dim))])
+    return [*product(*([[*range(-cutoff, cutoff+1)]]*dim))]
 
 def hk2tb_model(hk, hopping_vecs, ks=None):
     """