[PyTorch] Update RNG global states in tracker set_states (#2501)

set_all_rng_states in set_states
Signed-off-by: Robin Zhang <robinz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

transformer_engine/pytorch/distributed.py

@@ -818,6 +818,8 @@ class CudaRNGStatesTracker:
                A mapping from string names to RNG states.
         """
         self.states_ = states
+        # Update global states.
+        set_all_rng_states(self.states_)
 
     def add(self, name: str, seed: int) -> None:
         """