diff --git a/medcat-v2/medcat/components/linking/vector_context_model.py b/medcat-v2/medcat/components/linking/vector_context_model.py
index 296f820cc..9afac3f7f 100644
--- a/medcat-v2/medcat/components/linking/vector_context_model.py
+++ b/medcat-v2/medcat/components/linking/vector_context_model.py
@@ -344,7 +344,7 @@ def train(self, cui: str, entity: MutableEntity, doc: MutableDocument,
                 if name not in self.name2info:
                     continue
                 per_cui_status = self.name2info[name]['per_cui_status']
-                cui_status = per_cui_status[cui]
+                cui_status = per_cui_status.get(cui, None)
                 if cui_status == ST.PRIMARY_STATUS_NO_DISAMB:
                     # Set this name to always be disambiguated, even
                     # though it is primary
diff --git a/medcat-v2/medcat/trainer.py b/medcat-v2/medcat/trainer.py
index 231cae707..8f1f9f87d 100644
--- a/medcat-v2/medcat/trainer.py
+++ b/medcat-v2/medcat/trainer.py
@@ -438,7 +438,7 @@ def _train_supervised_for_project2(self,
                         cui=cui, name=ann['value'], mut_doc=mut_doc,
                         mut_entity=mut_entity, negative=deleted,
                         devalue_others=devalue_others)
-                except ValueError as ve:
+                except (ValueError, KeyError) as ve:
                     context_window = 20  # characters
                     splitter_left, splitter_right = "<", ">"
                     cur_text = doc['text']