added pre ad post process

54093417 · Mostofa Patwary · f9267205 · 54093417
Commit 54093417 authored May 12, 2021 by Mostofa Patwary
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

megatron/indexer.py megatron/indexer.py +2 -1

No files found.
--- a/megatron/indexer.py
+++ b/megatron/indexer.py
@@ -58,7 +58,8 @@ class IndexBuilder(object):
        #model = get_model(lambda: biencoder_model_provider(only_context_model \
        model = get_model(biencoder_model_provider(only_context_model \
            = only_context_model, biencoder_shared_query_context_model = \
-            self.biencoder_shared_query_context_model))
+            self.biencoder_shared_query_context_model,
+            pre_process=True, post_process=True))

        self.model = load_biencoder_checkpoint(model,
                only_context_model=only_context_model)