Update app.py

app.py

@@ -3,13 +3,21 @@ import torch
 import gym
 from models.attention_model_wrapper import Agent
 
+from wrappers.syncVectorEnvPomo import SyncVectorEnv
+from wrappers.recordWrapper import RecordEpisodeStatistics
+
+
+import matplotlib.pyplot as plt
+from matplotlib.collections import LineCollection
+from matplotlib.colors import ListedColormap, BoundaryNorm
+
+import gradio as gr
+
 device = "cpu"
 ckpt_path = "./runs/tsp-v0__ppo_or__1__1678160003/ckpt/12000.pt"
 agent = Agent(device=device, name="tsp").to(device)
 agent.load_state_dict(torch.load(ckpt_path, map_location=torch.device("cpu")))
 
-from wrappers.syncVectorEnvPomo import SyncVectorEnv
-from wrappers.recordWrapper import RecordEpisodeStatistics
 
 env_id = "tsp-v0"
 env_entry_point = "envs.tsp_vector_env:TSPVectorEnv"
@@ -75,9 +83,6 @@ default_data = np.array(
 
 # @title Helper function for plotting
 # colorline taken from https://nbviewer.org/github/dpsanders/matplotlib-examples/blob/master/colorline.ipynb
-import matplotlib.pyplot as plt
-from matplotlib.collections import LineCollection
-from matplotlib.colors import ListedColormap, BoundaryNorm
 
 
 def make_segments(x, y):
@@ -126,9 +131,6 @@ def plot(coords):
     return fig
 
 
-import gradio as gr
-
-
 def run_inference(data):
     data = data.astype(float).to_numpy()
     resulting_traj, final_return = inference(data)
@@ -149,4 +151,4 @@ demo = gr.Interface(
     ),
     [gr.Plot(label="Results Visualization"), gr.Code(label="Results", interactive=False)],
 )
-demo.launch()
+demo.launch(share=True)