Update main.py

main.py

@@ -83,7 +83,7 @@ async def execute_code(session_token: str, code: str):
         
         output = []
         while True:
-            msg = st.session_state.kernel_client.get_iopub_msg()
+            msg = kc.get_iopub_msg()
             if msg['msg_type'] == 'status' and msg['content']['execution_state'] == 'idle':
                 break
             elif msg['msg_type'] == 'error':