Update src/litbench_pipeline.py

src/litbench_pipeline.py

@@ -622,9 +622,16 @@ if __name__ == "__main__":
     random.seed(seed_no)
 
 
+    # Load model and tokenizer
+    bnb_config = BitsAndBytesConfig(
+        load_in_8bit=True,
+        bnb_8bit_use_double_quant=True,
+        bnb_8bit_quant_type="nf8",
+        bnb_8bit_compute_dtype=torch.bfloat16
+    )
     # Load model and tokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config)
     if model.device.type != 'cuda':
         model.to('cuda')