Update pages/2_CodeSimulator.py

pages/2_CodeSimulator.py

@@ -1,12 +1,26 @@
 import streamlit as st
 import torch
+import io
+import sys
 
-# Function to execute the input code
+# Function to execute the input code and capture print statements
 def execute_code(code):
+    # Redirect stdout to capture print statements
+    old_stdout = sys.stdout
+    sys.stdout = mystdout = io.StringIO()
+    
     global_vars = {}
     local_vars = {}
-    exec(code, global_vars, local_vars)
-    return local_vars
+    try:
+        exec(code, global_vars, local_vars)
+        output = mystdout.getvalue()
+    except Exception as e:
+        output = str(e)
+    finally:
+        # Reset redirect.
+        sys.stdout = old_stdout
+    
+    return output, local_vars
 
 st.title('PyTorch Code Runner')
 
@@ -15,13 +29,15 @@ code_input = st.text_area("Enter your PyTorch code here", height=300)
 
 # Button to execute the code
 if st.button("Run Code"):
-    try:
-        # Execute the code
-        output = execute_code(code_input)
-        
-        # Display the output
-        st.subheader('Output')
-        for key, value in output.items():
-            st.text(f"{key}: {value}")
-    except Exception as e:
-        st.error(f"Error: {e}")
+    # Execute the code and capture the output
+    output, variables = execute_code(code_input)
+    
+    # Display the output
+    st.subheader('Output')
+    st.text(output)
+    
+    # Display returned variables
+    if variables:
+        st.subheader('Variables')
+        for key, value in variables.items():
+            st.text(f"{key}: {value}")