import streamlit as st
import torch

# Function to execute the input code
def execute_code(code):
    global_vars = {}
    local_vars = {}
    exec(code, global_vars, local_vars)
    return local_vars

st.title('PyTorch Code Runner')

# Text area for inputting the PyTorch code
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}")