diff --git a/spacy/ml/models/coref.py b/spacy/ml/models/coref.py
index 24b5500a2..cfbe83a7a 100644
--- a/spacy/ml/models/coref.py
+++ b/spacy/ml/models/coref.py
@@ -22,9 +22,8 @@ def build_wl_coref_model(
     rough_k: int = 50,
     # TODO is this not a training loop setting?
     a_scoring_batch_size: int = 512,
-    # span predictor embeddings
-    sp_embedding_size: int = 64,
 ):
+    # TODO add model return types
     # TODO fix this
     try:
         dim = tok2vec.get_dim("nO")
diff --git a/spacy/ml/models/span_predictor.py b/spacy/ml/models/span_predictor.py
index b990b4019..c5cbb328c 100644
--- a/spacy/ml/models/span_predictor.py
+++ b/spacy/ml/models/span_predictor.py
@@ -17,6 +17,7 @@ def build_span_predictor(
     hidden_size: int = 1024,
     dist_emb_size: int = 64,
 ):
+    # TODO add model return types
     # TODO fix this
     try:
         dim = tok2vec.get_dim("nO")