diff --git a/quantmsrescore/__init__.py b/quantmsrescore/__init__.py
index 3b93d0b..27fdca4 100644
--- a/quantmsrescore/__init__.py
+++ b/quantmsrescore/__init__.py
@@ -1 +1 @@
-__version__ = "0.0.2"
+__version__ = "0.0.3"
diff --git a/quantmsrescore/ms2rescore.py b/quantmsrescore/ms2rescore.py
index 9d02c50..5251123 100644
--- a/quantmsrescore/ms2rescore.py
+++ b/quantmsrescore/ms2rescore.py
@@ -36,6 +36,7 @@ def __init__(self, filename: Union[Path, str], *args, **kwargs) -> None:
         self.user_params_metadata = self._get_userparams_metadata(self.peptide_ids[0].getHits()[0])
         self.rescoring_features = self._get_rescoring_features(self.peptide_ids[0].getHits()[0])
         self.skip_invalid_psm = 0
+        self.new_peptide_ids = []
 
     def __iter__(self) -> Iterable[PSM]:
         """
@@ -63,12 +64,16 @@ def __iter__(self) -> Iterable[PSM]:
 
         """
         for peptide_id in self.peptide_ids:
+            new_hits = []
             for peptide_hit in peptide_id.getHits():
                 psm = self._parse_psm(self.protein_ids, peptide_id, peptide_hit)
                 if psm is not None:
+                    new_hits.append(peptide_hit)
                     yield psm
                 else:
                     self.skip_invalid_psm += 1
+            peptide_id.setHits(new_hits)
+            self.new_peptide_ids.append(peptide_id)
 
     def _parse_psm(
             self,
@@ -239,12 +244,14 @@ def rescore_idxml(input_file, output_file, config) -> None:
         logging.warning(
             f"Removed {reader.skip_invalid_psm} PSMs without search engine features!"
         )
-
+        peptide_ids = reader.new_peptide_ids
+    else:
+        peptide_ids = reader.peptide_ids
     # Rescore
     rescore(config, psm_list)
 
     # Filter out PeptideHits within PeptideIdentification(s) that could not be processed by all feature generators
-    peptide_ids_filtered = filter_out_artifact_psms(psm_list, reader.peptide_ids)
+    peptide_ids_filtered = filter_out_artifact_psms(psm_list, peptide_ids)
 
     # Write
     writer = IdXMLWriter(output_file, reader.protein_ids, peptide_ids_filtered)