flopml
/

mingru

flpelerin commited on Nov 1, 2024

Commit

be0031f

1 Parent(s): de45ebc

update

Files changed (1) hide show

train.py CHANGED Viewed

@@ -168,7 +168,8 @@ for epoch in range(num_epochs):
             "step": step,
             "epoch": epoch + 1,
             "loss": loss.item(),
-            "average_hidden_state": avg_states
         })
         print(f"Epoch: {epoch + 1}/{num_epochs}, Step: {step}, Loss: {loss.item():.4f}, Hidden States: average = {avg_states}, variance = {var_states}")

             "step": step,
             "epoch": epoch + 1,
             "loss": loss.item(),
+            "average_hidden_state": avg_states,
+            "variance_hidden_state": var_states
         })
         print(f"Epoch: {epoch + 1}/{num_epochs}, Step: {step}, Loss: {loss.item():.4f}, Hidden States: average = {avg_states}, variance = {var_states}")