diff --git a/flagai/trainer.py b/flagai/trainer.py
index ea3e28c8..a72bc3db 100644
--- a/flagai/trainer.py
+++ b/flagai/trainer.py
@@ -707,8 +707,8 @@ def evaluate(self,
                 step_output = forward_step_func(data_iterator,
                                                 model,
                                                 mems=mems)
-                lm_loss, mems = step_output['loss'], step_output[
-                    'hidden_states']
+                lm_loss= step_output['loss']
+                # mems = step_output['hidden_states']
                 '''when contiguous memory optimizations are enabled, the buffers
                 allocated by the optimizations are deallocated during backward pass
                 in the absence of backward pass the buffers should be reset after each