test

train.py

@@ -155,6 +155,7 @@ for epoch in range(num_epochs):
         # Reset hidden states if needed
         h_states = h_states if (step % reset_state_every!= 0) else None
         avg_states = sum([torch.mean(h_states[i]).item() for i in range(len(h_states))]) / len(h_states) if h_states is not None else None
+        var_states = torch.var(torch.cat(hidden_states_list, dim=0)).item() if hidden_states_list else None
 
         optimizer.zero_grad()
         _, h_states, loss = model.forward(input_ids, h_states)