Update app.py

app.py

@@ -44,9 +44,9 @@ def Compute_Parameters_attention(hidden_size, kv_hidden_size, is_bias, act_func,
     # attention: 
     # layernorm: h/2h
     if act_func == "LLaMA":
-        num_parameters_mlp = hidden_size # RMSNorm
+        num_parameters_attention = hidden_size # RMSNorm
     else:
-        num_parameters_mlp = 2 * hidden_size # LayerNorm
+        num_parameters_attention = 2 * hidden_size # LayerNorm
     # QKV weight: 3h*h/tp, bias: 3h/tp
     # output linear weight: h*h/tp, bias: h
     num_parameters_attention_Q_weight = hidden_size * hidden_size / tp
@@ -85,7 +85,7 @@ def Compute_Parameters(seq_length, vocab_size, layer_num, hidden_size, ffn_size,
     kv_hidden_size = hidden_size / head_num * group_query_num
     
     # input part
-    num_parameters_input = Compute_Parameters_input(seq_length, hidden_size, vocab_size, tp)
+    num_parameters_input = Compute_Parameters_input(seq_length, hidden_size, vocab_size, act_func, tp)
 
     # middle layers part
     num_parameters_attention = Compute_Parameters_attention(hidden_size, kv_hidden_size, is_bias, act_func, tp)