Update app.py

app.py

@@ -12,12 +12,12 @@ def plot_forecast(num_param, batch_size, precision, seq_len):
     # Model Parameters: N×precision
     y1 = num_param * precision / (1024**3)
 
-    # Optimizer States: 2×N×precision
-    y2 = 2 * num_param * precision / (1024**3)
-
     # Activations: B×Sequence Length×K×precision
     K = 4.6894e-04 * num_param + 1.8494e06
-    y3 = batch_size * seq_len * K * precision / (1024**3)
+    y2 = batch_size * seq_len * K * precision / (1024**3)
+
+    # Optimizer States: 2×N×precision
+    y3 = 2 * num_param * precision / (1024**3)
 
     # Gradients: N×precision
     y4 = num_param * 4 / (1024**3)
@@ -33,8 +33,8 @@ def plot_forecast(num_param, batch_size, precision, seq_len):
 
     # Add text labels inside the bars
     ax.text(0, y1 / 2, "Model Parameters", ha="center", va="center", color="white", fontweight="bold")
-    ax.text(0, y1 + y2 / 2, "Optimizer States", ha="center", va="center", color="white", fontweight="bold")
-    ax.text(0, y1 + y2 + y3 / 2, "Activations", ha="center", va="center", color="white", fontweight="bold")
+    ax.text(0, y1 + y2 / 2,"Activations", ha="center", va="center", color="white", fontweight="bold")
+    ax.text(0, y1 + y2 + y3 / 2, "Optimizer States", ha="center", va="center", color="white", fontweight="bold")
     ax.text(0, y1 + y2 + y3 + y4 / 2, "Gradients", ha="center", va="center", color="white", fontweight="bold")
 
     # remove x axis