diff --git a/rf2aa/data/data_loader_utils.py b/rf2aa/data/data_loader_utils.py
index bd66579..0b3c38a 100644
--- a/rf2aa/data/data_loader_utils.py
+++ b/rf2aa/data/data_loader_utils.py
@@ -455,7 +455,7 @@ def merge_a3m_homo(msa_orig, ins_orig, nmer, mode="default"):
              ins[N:, start:start+L] = ins_orig[1:]
              start += L        
 
-     return msa, ins
+     return {"msa": msa, "ins": ins}
 
 def merge_msas(a3m_list, L_s):
     """
diff --git a/rf2aa/data/merge_inputs.py b/rf2aa/data/merge_inputs.py
index bf12ff8..7917fa1 100644
--- a/rf2aa/data/merge_inputs.py
+++ b/rf2aa/data/merge_inputs.py
@@ -43,8 +43,8 @@ def merge_protein_inputs(protein_inputs, deterministic: bool = False):
 
         if len(unique_a3m) >1:
             a3m_out = unique_a3m[0]
-            for i in range(1, len(unique_a3m)):
-                a3m_out = join_msas_by_taxid(a3m_out, a3m_list[i])
+            for unq_a3m in unique_a3m[1:]:
+                a3m_out = join_msas_by_taxid(a3m_out, unq_a3m)
             a3m_out = expand_multi_msa(a3m_out, unique_hashes, hash_list, unique_lengths_list, lengths_list)
         else:
             a3m  = unique_a3m[0]