From eb5820b5933d60f8c5b854a7717a1db1b4c9cb2c Mon Sep 17 00:00:00 2001 From: Paul O'Leary McCann Date: Mon, 5 Jul 2021 21:08:42 +0900 Subject: [PATCH] Improve take_vecs implementation This pulls out references to needed bits so that other parts (the larger embeddings) can be freed before backprop. --- spacy/ml/models/coref.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/spacy/ml/models/coref.py b/spacy/ml/models/coref.py index f83014344..31643d248 100644 --- a/spacy/ml/models/coref.py +++ b/spacy/ml/models/coref.py @@ -356,9 +356,11 @@ def build_take_vecs() -> Model[SpanEmbeddings, Floats2d]: def take_vecs_forward(model, inputs: SpanEmbeddings, is_train) -> Floats2d: + idxs = inputs.indices + lens = inputs.vectors.lengths def backprop(dY: Floats2d) -> SpanEmbeddings: - vecs = Ragged(dY, inputs.vectors.lengths) - return SpanEmbeddings(inputs.indices, vecs) + vecs = Ragged(dY, lens) + return SpanEmbeddings(idxs, vecs) return inputs.vectors.data, backprop