def test_safetensors_load_from_local(self):
        """
        This test checks that we can load safetensors from a checkpoint that only has those on the Hub
        """
        with tempfile.TemporaryDirectory() as tmp:
            location = snapshot_download("hf-internal-testing/tiny-bert-flax-only", cache_dir=tmp)
            flax_model = FlaxBertModel.from_pretrained(location)

        with tempfile.TemporaryDirectory() as tmp:
            location = snapshot_download("hf-internal-testing/tiny-bert-flax-safetensors-only", cache_dir=tmp)
            safetensors_model = FlaxBertModel.from_pretrained(location)

        self.assertTrue(check_models_equal(flax_model, safetensors_model))

        This test checks that we can load safetensors from a checkpoint that only has those on the Hub
        
dpkd/transformers/tests/test_modeling_flax_utils.py