fix tests

Author: sayakpaul

tests/models/unets/test_models_unet_2d_condition.py

@@ -1086,16 +1086,19 @@ def test_load_sharded_checkpoint_device_map_from_hub_local_subfolder(self):
 
     def test_wrong_device_map_raises_error(self):
         with self.assertRaises(ValueError) as err_ctx:
-            _ = self.model_class.from_pretrained("hf-internal-testing/unet2d-sharded-dummy-subfolder", device_map=-1)
-            msg_substring = "You can't pass device_map as a negative int"
-            assert msg_substring in str(err_ctx.exception)
+            _ = self.model_class.from_pretrained(
+                "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map=-1
+            )
+
+        msg_substring = "You can't pass device_map as a negative int"
+        assert msg_substring in str(err_ctx.exception)
 
     @require_torch_gpu
     @parameterized.expand([0, "cuda", torch.device("cuda"), torch.device("cuda:0")])
     def test_passing_non_dict_device_map_works(self, device_map):
         _, inputs_dict = self.prepare_init_args_and_inputs_for_common()
         loaded_model = self.model_class.from_pretrained(
-            "hf-internal-testing/unet2d-sharded-dummy-subfolder", device_map=device_map
+            "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map=device_map
         )
         output = loaded_model(**inputs_dict)
         assert output.sample.shape == (4, 4, 16, 16)