Update app.py

app.py

@@ -12,7 +12,7 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 base_model_name = "google/gemma-7b"
 adapter_model_name = "samidh/cope-g7bq-2c-hs.s1.5fpc.9-sx.s1.5.9o-VL.s1.5.9-HR.s5-SH.s5-l5e5-e3-d25-r8"
 
-model = AutoModelForCausalLM.from_pretrained(base_model_name, token=os.environ['HF_TOKEN'])
+model = AutoModelForCausalLM.from_pretrained(base_model_name, token=os.environ['HF_TOKEN'], device_map="auto")
 model = PeftModel.from_pretrained(model, adapter_model_name, token=os.environ['HF_TOKEN'])
 model.merge_and_unload()
 
@@ -86,6 +86,7 @@ def predict(content, policy):
     with torch.no_grad():
         outputs = model(input_ids)
         logits = outputs.logits[:, -1, :]  # Get logits for the last token
+        model.gradient_checkpointing_enable()
         predicted_token_id = torch.argmax(logits, dim=-1).item()
     decoded_output = tokenizer.decode([predicted_token_id])
     if decoded_output == '0':