Update app.py

add CUDA config

app.py

@@ -1,10 +1,18 @@
 import gradio as gr
+import torch
 from transformers import AutoModel
 
 @spaces.GPU
 def get_model_summary(model_name):
-    model = AutoModel.from_pretrained(model_name)
+    # Check if CUDA is available and set the device accordingly
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Load the model and move it to the selected device
+    model = AutoModel.from_pretrained(model_name).to(device)
+    
+    # Return the model's architecture as a string
     return str(model)
 
+# Create the Gradio interface
 interface = gr.Interface(fn=get_model_summary, inputs="text", outputs="text")
 interface.launch()