Update app.py

app.py

@@ -77,8 +77,10 @@ def plot_pens(tflpos_card, utilization, num_gps, training_days):
     plt.axvline(ns[best], color='red')
     plt.xlabel('model size')
     plt.ylabel('loss')
+    fig.savefig("/tmp/tmp.jpg")
+    plt.close()
 
-    return fig, c, round(losses[best], 3), best_model_size, best_dataset_size
+    return "/tmp/tmp.jpg", c, round(losses[best], 3), best_model_size, best_dataset_size
 
 
 if __name__ == "__main__":
@@ -86,12 +88,12 @@ if __name__ == "__main__":
         fn=plot_pens,
         inputs=[
             gr.Textbox(label="TFLOP/s pre Card",value="40"),
-            gr.Slider(label="System Utilization", minimum=0, maximum=1, step=0.01,value=0.25),
-            gr.Textbox(label="Number of cards",value="1"),
-            gr.Textbox(label="Training Days",value="7")
+            gr.Slider(label="GPU Utilization", minimum=0, maximum=1, step=0.01,value=0.25),
+            gr.Textbox(label="Number of cards"),
+            gr.Textbox(label="Training Days")
         ],
         outputs=[
-            gr.Plot(label="Estimated Loss"),
+            gr.Image(label="Estimated Loss"),
             gr.Label(label="Total Compute Budget"),
             gr.Label(label="Estimated Final Loss"),
             gr.Label(label="Optimal Model Size"),
@@ -100,5 +102,5 @@ if __name__ == "__main__":
         title="Compute-Optimal Model Estimator",
         description=description,
         article=article,
-        live=False
+        live=True
     ).launch()