fix

train.py

@@ -154,7 +154,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 = torch.mean(h_states).item() if h_states is not None 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
 
         optimizer.zero_grad()
         _, h_states, loss = model.forward(input_ids, h_states)