Add num_gpus

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import matplotlib.pyplot as plt
 
 
-def plot_forecast(num_param, precision, grad_ckpt, batch_size, seq_len):
+def plot_forecast(num_param, precision, grad_ckpt, batch_size, seq_len, num_gpus):
     # Convert number (input as B)
     num_param = float(num_param) * 1e9
 
@@ -100,6 +100,8 @@ with gr.Blocks() as demo:
                 seq_len = gr.Slider(1, 1000, label="Sequence Length", step=1, value=256)
 
             with gr.Accordion("Advanced", open=False):
+                with gr.Accordion("Hardware"):
+                    num_gpus = gr.Number(1, label="Number of GPUs")
                 with gr.Accordion("Data"):
                     grad_ckpt = gr.Checkbox(False, label="Gradient Checkpointing")
 
@@ -108,7 +110,7 @@ with gr.Blocks() as demo:
         with gr.Column():
             plot = gr.Plot(label="forecast", format="png")
 
-    submit.click(plot_forecast, [num_param, precision, grad_ckpt, batch_size, seq_len], plot)
+    submit.click(plot_forecast, [num_param, precision, grad_ckpt, batch_size, seq_len, num_gpus], plot)
 
 if __name__ == "__main__":
     demo.launch()