diff --git a/ocaml/Input_determinants.ml b/ocaml/Input_determinants.ml
index bc476c12..6c47d966 100644
--- a/ocaml/Input_determinants.ml
+++ b/ocaml/Input_determinants.ml
@@ -131,9 +131,11 @@ end = struct
     |> States_number.of_int
   ;;
 
-  let write_n_states_diag n =
-    States_number.to_int n
-    |> Ezfio.set_determinants_n_states_diag
+  let write_n_states_diag ~n_states n =
+    let n_states = States_number.to_int n_states
+    and n = States_number.to_int n
+    in
+    Ezfio.set_determinants_n_states_diag (max n_states n)
   ;;
 
 
@@ -232,20 +234,27 @@ end = struct
 
   let read_psi_coef () =
     if not (Ezfio.has_determinants_psi_coef ()) then
-        Ezfio.ezfio_array_of_list ~rank:1 ~dim:[| 1 |] ~data:[1.]
-        |> Ezfio.set_determinants_psi_coef 
-      ;
+      begin
+        let n_states =
+          read_n_states ()
+          |> States_number.to_int
+        in
+        Ezfio.ezfio_array_of_list ~rank:2 ~dim:[| 1 ; n_states |] 
+          ~data:(List.init n_states ~f:(fun i -> if (i=0) then 1. else 0. ))
+          |> Ezfio.set_determinants_psi_coef 
+      end;
     Ezfio.get_determinants_psi_coef ()
     |> Ezfio.flattened_ezfio
     |> Array.map ~f:Det_coef.of_float
   ;;
 
-  let write_psi_coef ~n_det c =
+  let write_psi_coef ~n_det ~n_states c =
     let n_det = Det_number.to_int n_det
     and c = Array.to_list c
             |> List.map ~f:Det_coef.to_float
+    and n_states = States_number.to_int n_states
     in
-    Ezfio.ezfio_array_of_list ~rank:1 ~dim:[| n_det |] ~data:c
+    Ezfio.ezfio_array_of_list ~rank:2 ~dim:[| n_det ; n_states |] ~data:c
     |> Ezfio.set_determinants_psi_coef 
   ;;
 
@@ -339,14 +348,14 @@ end = struct
      write_mo_label mo_label;
      write_n_det n_det;
      write_n_states n_states;
-     write_n_states_diag n_states_diag;
+     write_n_states_diag ~n_states:n_states n_states_diag;
      write_n_det_max_jacobi n_det_max_jacobi;
      write_threshold_generators threshold_generators;
      write_threshold_selectors threshold_selectors;
      write_read_wf read_wf;
      write_expected_s2 expected_s2;
      write_s2_eig s2_eig;
-     write_psi_coef ~n_det:n_det psi_coef;
+     write_psi_coef ~n_det:n_det psi_coef ~n_states:n_states;
      write_psi_det ~n_int:n_int ~n_det:n_det psi_det;
   ;;