Update app.py

app.py

@@ -1,9 +1,5 @@
 import streamlit as st
 
-def number_field(label, c1, c2, **kwargs):
-  c1.write(label)
-  return c2.number_input('', **kwargs)
-
 def calc_exec_time(comp_flop, mem_bytes):
   return (comp_flop/TFLOPS + mem_bytes/GB_S)*1000
   
@@ -31,12 +27,14 @@ GB_S = 1935e9
 st.sidebar.header("Transformer parameters")
 col1, col2 = st.sidebar.columns([2, 4])
 
-bs = number_field('Batch size', col1, col2, value=10)
-h = number_field('Num heads', col1, col2, value=16)
-d = number_field('Dimension', col1, col2, value=768)
-n_start = number_field('Start seq', col1, col2, value=1)
-n = number_field('End seq', col1, col2, value=1024)
-l = number_field('Num layers', col1, col2, value=24)
+bs = st.sidebar.number_input('Batch size', value=10)
+h = st.sidebar.number_input('Num heads',value=16)
+d = st.sidebar.number_input('Dimension', value=768)
+l = st.sidebar.number_input('Num layers', value=24)
+n_start = st.sidebar.number_input('Start seq', value=1)
+n = st.sidebar.number_input('End seq', value=1024)
+
+st.sidebar.header("GPU parameters")
 
 st.header('Attention layer')