Update app.py

app.py

@@ -63,16 +63,16 @@ cls_model = cls_model.to(device)
 
 @spaces.GPU(duration=1500)  # maximum duration allowed during startup
 def compile_transformer():
-    with spaces.aoti_capture(cls_model) as call:
+    with spaces.aoti_capture(chat_model.model) as call:
         chat_model("arbitrary example prompt")
 
-    exported = torch.export.export(cls_model, args=call.args, kwargs=call.kwargs)
+    exported = torch.export.export(chat_model.model, args=call.args, kwargs=call.kwargs)
     return spaces.aoti_compile(exported)
 
 print("\n=== Model Compilation ===")
 
 compiled_transformer = compile_transformer()
-spaces.aoti_apply(compiled_transformer, cls_model)
+spaces.aoti_apply(compiled_transformer, chat_model.model)
 
 
 def generate_responses(model, tokenizer, prompts):