From 07794179037f8bcbf061902bcc87dd8ef43359e7 Mon Sep 17 00:00:00 2001
From: Juarez Bochi <juarez.bochi@grammarly.com>
Date: Mon, 18 Dec 2023 11:19:44 -0500
Subject: [PATCH] Fix --encode-only

---
 t5/t5.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/t5/t5.py b/t5/t5.py
index d5ccdd9c..65c758e4 100644
--- a/t5/t5.py
+++ b/t5/t5.py
@@ -394,8 +394,7 @@ if __name__ == "__main__":
     if args.encode_only:
         print("[INFO] Encoding with T5...", flush=True)
         print(args.prompt, flush=True)
-        embeddings = model.wte(prompt)
-        encoder_output = model.encoder(embeddings, mask=None)
+        encoder_output = model.encode(prompt)
         print(encoder_output, flush=True)
         exit(0)