Update app.py

app.py

@@ -48,7 +48,13 @@ async def run_code(code: Code):
     # img_buffer.seek(0)  # Reset buffer position
 
     file_path = "graph.pdf"
-    fig.save_fig(file_path)
+
+    if code.code == "plt.subplots(":
+        
+        exec(code.code + "\nfig.savefig(file_path)")
+    else:
+        exec(code.code + "\nplt.savefig(file_path)")
+
     plt.close()
     # plt.savefig(file_path)
     # plt.close()