fix bug

Author: Your Name

paddlenlp/trl/embedding_trainer.py

@@ -77,7 +77,7 @@ def get_current_rng_state(self):
         return {
             "cpu": [paddle.framework.core.default_cpu_generator().get_state()],
             "cuda": [paddle.get_rng_state()],
-            "hybrid": [fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()],
+            "hybrid": [fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()] if self.args.use_hybrid_parallel else []
         }
 
     def reset_rng_state(self, states, index=0):
@@ -86,13 +86,13 @@ def reset_rng_state(self, states, index=0):
             raise ValueError("The length of state should be 3")
         cpu_state = states["cpu"][index]
         cuda_state = states["cuda"][index]
-        hybrid_state = states["hybrid"][index]
         paddle.framework.core.default_cpu_generator().set_state(cpu_state)
         # TODO(daisiming): support xpu and other custom devices.
         if core.is_compiled_with_cuda():
             for j in range(core.get_cuda_device_count()):
                 core.default_cuda_generator(j).set_state(cuda_state[j])
         if self.args.use_hybrid_parallel:
+            hybrid_state = states["hybrid"][index]
             fleet.meta_parallel.get_rng_state_tracker().set_states_tracker(hybrid_state)
 
     def accum_forward_backward(self, model):