Update train.py

train.py

@@ -89,7 +89,7 @@ for epoch in range(num_epochs):
         #    h_states = None
             
         h_states = h_states if (i / batch_size) % reset_state_every != 0 else None
-        str_states = ''.join(['{:.3f}, '.format(h_states[0][0][0][i].item()) for i in range(10)]) if h_states is not None else 'None'
+        str_states = ''.join(['{:.3f}, '.format(h_states[0][0][i].item()) for i in range(10)]) if h_states is not None else 'None'
 
         optimizer.zero_grad()
         _, h_states, loss = model.forward(input_ids, h_states)