init

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import spaces
 import torch
 from transformers import AutoTokenizer
 from sarm_llama import LlamaSARM
@@ -18,7 +19,6 @@ model = LlamaSARM.from_pretrained(
     sae_latent_size=65536,
     sae_k=192,
     device_map="auto",  # <<< KEY CHANGE HERE
-    trust_remote_code=True, 
     torch_dtype=torch.bfloat16
 )
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
@@ -28,7 +28,7 @@ DEVICE = model.device
 print(f"Model loaded successfully on device: {DEVICE}")
 
 # --- 2. Define the Inference Function ---
-
+@spaces.GPU
 def get_reward_score(prompt: str, response: str) -> float:
     """
     Receives a prompt and a response, and returns the reward score calculated by the SARM model.