diff --git a/barrOs_library.py b/barrOs_library.py
index 5f92519054d93f88bcf8d74f0321c10927da2ea8..98416952c5bd9175ea0dcb67630cde82f46b5c63 100644
--- a/barrOs_library.py
+++ b/barrOs_library.py
@@ -2302,9 +2302,9 @@ def run_barros(arguments, offset = 1, step = 2, local_angle_threshold = 25, dist
             target_chains = []
             if len(chains_inpdb) > 1:
                 for chainID in chains_inpdb:
-                    if input_mode == 'pdb_file':
+                    if '.pdb' in pdbID or '.ent' in pdbID:
                         chain_pdb, _ = extract_chain(pdb_file, chain, outpdb = 'downloaded_pdbs/{}_{}.pdb'.format(pdbID, chainID))
-                        target_chains.append([pdbID, chainID, chain_pdb, protein_type, membrane_thickness])
+                        target_chains.append([pdbID.split('.')[0], chainID, chain_pdb, protein_type, membrane_thickness])
                     else:  
                         chain_pdb, protein_type, membrane_thickness = download_pdb('{}_{}'.format(pdbID.split('_')[0], chainID))
                         target_chains.append([pdbID.split('_')[0], chainID, chain_pdb, protein_type, membrane_thickness])