Updated lines 166-177

app.py

@@ -168,12 +168,15 @@ def plot_trellis_with_path(trellis, path):
     trellis_with_path = trellis.clone()
     for _, p in enumerate(path):
         trellis_with_path[p.time_index, p.token_index] = float("nan")
-    plt.imshow(trellis_with_path.T, origin="lower")
-    plt.title("The path found by backtracking")
-    plt.tight_layout()
-    return plt 
+    fig, ax = plt.subplots()
+    ax.imshow(trellis_with_path.T, origin="lower")
+    ax.set_title("The path found by backtracking")
+    ax.set_xlabel("Time")
+    ax.set_ylabel("Labels")
+    fig.tight_layout()
+    return fig
     
-st.pyplot(plt)
+st.pyplot(fig)
 
 # Part J: Merge Repeats | Segments
 # Merge the labels