[Flax] Update no init test for Flax v0.7.1 (#28735)

tests/test_modeling_flax_common.py

@@ -984,7 +984,7 @@ def test_no_automatic_init(self):
 
             # Check if we params can be properly initialized when calling init_weights
             params = model.init_weights(model.key, model.input_shape)
-            self.assertIsInstance(params, FrozenDict)
+            assert isinstance(params, (dict, FrozenDict)), f"params are not an instance of {FrozenDict}"
             # Check if all required parmas are initialized
             keys = set(flatten_dict(unfreeze(params)).keys())
             self.assertTrue(all(k in keys for k in model.required_params))