fixed CUDA issue for HF spaces

app.py

@@ -1,26 +1,30 @@
 import gradio as gr
 from generate import generate_text
+import torch
 
 def generate(prompt):
-    # Redirect print output to capture the generated text
-    import io
-    import sys
-    old_stdout = sys.stdout
-    new_stdout = io.StringIO()
-    sys.stdout = new_stdout
-    
-    # Generate text with default parameters
-    generate_text(
-        prompt=prompt,
-        max_length=100,  # default value
-        num_sequences=5  # default value
-    )
-    
-    # Get the output and restore stdout
-    output = new_stdout.getvalue()
-    sys.stdout = old_stdout
-    
-    return output
+    try:
+        # Redirect print output to capture the generated text
+        import io
+        import sys
+        old_stdout = sys.stdout
+        new_stdout = io.StringIO()
+        sys.stdout = new_stdout
+        
+        # Generate text with default parameters
+        generate_text(
+            prompt=prompt,
+            max_length=100,  # default value
+            num_sequences=5  # default value
+        )
+        
+        # Get the output and restore stdout
+        output = new_stdout.getvalue()
+        sys.stdout = old_stdout
+        
+        return output
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
 
 # Create the Gradio interface
 demo = gr.Interface(
@@ -32,4 +36,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch(share=False)  # set share=True if you want to generate a public URL 
+    demo.launch() 

generate.py

@@ -9,7 +9,7 @@ print(f"Using device: {device}")
 
 # Initialize model and load trained weights
 model = GPT(GPTConfig())
-model.load_state_dict(torch.load('best_model.pt'))
+model.load_state_dict(torch.load('best_model.pt', map_location=torch.device(device)))
 model.to(device)
 model.eval()  # Set to evaluation mode