diff --git a/jass_preprocessing/compute_score.py b/jass_preprocessing/compute_score.py
index 7e026c0ef2e9448eecc38280851a94080eb7a69f..1a1217ac39acadc3552c68134cc57c6d86317512 100644
--- a/jass_preprocessing/compute_score.py
+++ b/jass_preprocessing/compute_score.py
@@ -30,7 +30,7 @@ def compute_z_score(mgwas):
         
         mgwas.loc[mgwas["computed_z"].isin([np.inf]), 'computed_z'] = mgwas["z"]
 
-    elif 'beta' in mgwas.columns:
+    elif 'beta_or_Z' in mgwas.columns:
         sign_vect = np.sign(mgwas.beta)
         mgwas["computed_z"] = mgwas["computed_z"].replace(np.inf, 37.537836095576054)
     else: