Fix --encode-only

This commit is contained in:
Juarez Bochi 2023-12-18 11:19:44 -05:00
parent 83b68a5bdb
commit 0779417903
No known key found for this signature in database
GPG Key ID: 34CCBB77DC8BEBB6

View File

@ -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)