From 508f811b93035f076e18778fe08106f15abfa8a7 Mon Sep 17 00:00:00 2001
From: Juarez Bochi <juarez.bochi@grammarly.com>
Date: Mon, 6 Nov 2023 23:35:37 -0500
Subject: Add support for MADLAD400 (#1285)

* Add support for madlad

* Add support for quantized MADLAD
---
 candle-examples/examples/quantized-t5/main.rs | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

(limited to 'candle-examples/examples/quantized-t5/main.rs')

diff --git a/candle-examples/examples/quantized-t5/main.rs b/candle-examples/examples/quantized-t5/main.rs
index 5a1cdf0c..0ea2e0bd 100644
--- a/candle-examples/examples/quantized-t5/main.rs
+++ b/candle-examples/examples/quantized-t5/main.rs
@@ -173,7 +173,11 @@ fn main() -> Result<()> {
         .to_vec();
     let input_token_ids = Tensor::new(&tokens[..], device)?.unsqueeze(0)?;
     let mut model = builder.build_model()?;
-    let mut output_token_ids = [builder.config.pad_token_id as u32].to_vec();
+    let mut output_token_ids = [builder
+        .config
+        .decoder_start_token_id
+        .unwrap_or(builder.config.pad_token_id) as u32]
+    .to_vec();
     let temperature = if args.temperature <= 0. {
         None
     } else {
-- 
cgit v1.2.3