[Tests] fix more sharding tests (#8797)

* fix

* fix

* ugly

* okay

* fix more

* fix oops

Author: sayakpaul

tests/models/test_modeling_common.py

@@ -885,11 +885,11 @@ def test_model_parallelism(self):
 
     @require_torch_gpu
     def test_sharded_checkpoints(self):
+        torch.manual_seed(0)
         config, inputs_dict = self.prepare_init_args_and_inputs_for_common()
         model = self.model_class(**config).eval()
         model = model.to(torch_device)
 
-        torch.manual_seed(0)
         base_output = model(**inputs_dict)
 
         model_size = compute_module_sizes(model)[""]
@@ -909,7 +909,8 @@ def test_sharded_checkpoints(self):
             new_model = new_model.to(torch_device)
 
             torch.manual_seed(0)
-            _, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+            if "generator" in inputs_dict:
+                _, inputs_dict = self.prepare_init_args_and_inputs_for_common()
             new_output = new_model(**inputs_dict)
 
             self.assertTrue(torch.allclose(base_output[0], new_output[0], atol=1e-5))
@@ -942,7 +943,8 @@ def test_sharded_checkpoints_device_map(self):
             new_model = new_model.to(torch_device)
 
             torch.manual_seed(0)
-            _, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+            if "generator" in inputs_dict:
+                _, inputs_dict = self.prepare_init_args_and_inputs_for_common()
             new_output = new_model(**inputs_dict)
             self.assertTrue(torch.allclose(base_output[0], new_output[0], atol=1e-5))