Update pages/2_CodeSimulator.py

pages/2_CodeSimulator.py

@@ -1,5 +1,4 @@
 import streamlit as st
-import torch
 import io
 import sys
 
@@ -9,7 +8,7 @@ def execute_code(code):
     old_stdout = sys.stdout
     sys.stdout = mystdout = io.StringIO()
     
-    global_vars = {}
+    global_vars = {"torch": torch}
     local_vars = {}
     try:
         exec(code, global_vars, local_vars)
@@ -40,8 +39,11 @@ print("Broadcast Multiplication:\\n", tensor_broadcast_mul)
 
 # Button to execute the code
 if st.button("Run Code"):
+    # Prepend the import statement
+    code_to_run = "import torch\n" + code_input
+    
     # Execute the code and capture the output
-    output, variables = execute_code(code_input)
+    output, variables = execute_code(code_to_run)
     
     # Display the output
     st.subheader('Output')