Update app.py

app.py

@@ -37,8 +37,8 @@ if __name__ == "__main__":
                 **model_params
             )
         
-        plt.axis('off')
         fig, (ax1, ax2) = plt.subplots(1, 2)
+        fig.axis('off')
         ax1.imshow(preds[0])
         ax2.imshow(flows[0][0])
         st.pyplot(fig)