gated model updates

app.py

@@ -26,11 +26,13 @@ lm_model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
-# Load the reward model
+# Load the reward model - fix the offloading issue
+print("Loading reward model...")
 RM = AutoModelForCausalLMWithValueHead.from_pretrained(
     'ray24724919/plan2align_rm',
     torch_dtype=torch_dtype,
-    device_map="auto"
+    device_map={"": 0},  # Force model to stay on GPU (device 0)
+    offload_folder=None,  # Disable offloading
 )
 RM.eval()
 print("Models loaded successfully!")