update

train.py

@@ -170,7 +170,7 @@ for epoch in range(num_epochs):
             "loss": loss.item(),
             "average_hidden_state": avg_states
         })
-        print(f"Epoch: {epoch + 1}/{num_epochs}, Step: {step}, Loss: {loss.item():.4f}, Average Hidden State: {avg_states}, Hidden State Variance: {var_states}")
+        print(f"Epoch: {epoch + 1}/{num_epochs}, Step: {step}, Loss: {loss.item():.4f}, Hidden States: average = {avg_states}, variance = {var_states}")
 
         # Perform validation at specified intervals
         if step % validate_every == 0: