wrap automodel for zerogpu

langchainapp.py

@@ -43,8 +43,13 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 ### Utils  
 hf_token, yi_token = load_env_variables()  
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token, trust_remote_code=True)  
-nvidiamodel = AutoModel.from_pretrained(model_name, token=hf_token, trust_remote_code=True).to(device)  
+@spaces.GPU  
+def load_model():  
+    tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token, trust_remote_code=True)  
+    return AutoModel.from_pretrained(model_name, token=hf_token, trust_remote_code=True).to(device)  
+  
+# Load model  
+nvidiamodel = load_model()
 # nvidiamodel.set_pooling_include_prompt(include_prompt=False) 
 
 def clear_cuda_cache():