Repository: huggingface/transformers
File: tests/models/dab_detr/test_modeling_dab_detr.py (Python)

Code before:
    def setUp(self):
        self.model_tester = DabDetrModelTester(self)
        self.config_tester = ConfigTester(self, config_class=DabDetrConfig, has_text_modality=False)

    def test_config(self):
        self.config_tester.run_common_tests()

    def test_dab_detr_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_dab_detr_model(*config_and_inputs)

    def test_dab_detr_object_detection_head_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_dab_detr_object_detection_head_model(*config_and_inputs)

    def test_load_save_without_tied_weights(self):
        # DabDetrForObjectDetection forces `bbox_embed` to be tied by `self.x = y`
        # Run only DabDetrModel by overriding
        for model_class in [DabDetrModel]:
            config, _ = self.model_tester.prepare_config_and_inputs_for_common()
            config.tie_word_embeddings = False
            try:
                config.get_text_config().tie_word_embeddings = False
            except Exception as e:
                # Config may not have a text config
                print(f"Could not set text config tie_word_embeddings: {e}")

            model = model_class(config)  # we init the model without tie
            # if this test fails later on, it means init tied the weights
            with tempfile.TemporaryDirectory() as d:
                model.save_pretrained(d)
                with safe_open(f"{d}/model.safetensors", framework="pt") as f:
                    serialized_keys = f.keys()

                    model_reloaded, infos = model_class.from_pretrained(d, output_loading_info=True)
                    # Checking the state dicts are correct

                    reloaded_state = model_reloaded.state_dict()
                    for k, v in model.state_dict().items():
                        with self.subTest(k):
                            torch.testing.assert_close(
                                v,
                                reloaded_state[k],
                                msg=lambda x: f"{model_class.__name__}: Tensor {k}: {x}. Key {k} was serialized: {k in serialized_keys}. If `False`, this means it was probably aliased and safetensors removed it. If `True` it means `_init_weights` overwrote that key",
                            )

                # Checking there was no complain of missing weights
                self.assertEqual(
                    infos["missing_keys"],
                    set(),
                    "Given that the loaded weights are the same, the issue is in `tie_weights`: it tied these keys and removed them from serialization. But because of tiying (hardcoded or not) the previous check is fine.\

Review comment (suggestion):
I feel like this could be removed altogether (or adjusted properly). We are overriding and only check one class so we can specify it here

Code after:
    def setUp(self):
        self.model_tester = DabDetrModelTester(self)
        self.config_tester = ConfigTester(self, config_class=DabDetrConfig, has_text_modality=False)

    def test_config(self):
        self.config_tester.run_common_tests()

    def test_dab_detr_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_dab_detr_model(*config_and_inputs)

    def test_dab_detr_object_detection_head_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_dab_detr_object_detection_head_model(*config_and_inputs)

    def test_load_save_without_tied_weights(self):
        # DabDetrForObjectDetection forces `bbox_embed` to be tied by `self.x = y`
        # Run only DabDetrModel by overriding
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        config.tie_word_embeddings = False
        config.get_text_config().tie_word_embeddings = False

        model = DabDetrModel(config)  # we init the model without tie
        # if this test fails later on, it means init tied the weights
        with tempfile.TemporaryDirectory() as d:
            model.save_pretrained(d)
            with safe_open(f"{d}/model.safetensors", framework="pt") as f:
                serialized_keys = f.keys()

                model_reloaded, infos = DabDetrModel.from_pretrained(d, output_loading_info=True)
                # Checking the state dicts are correct

                reloaded_state = model_reloaded.state_dict()
                for k, v in model.state_dict().items():
                    with self.subTest(k):
                        torch.testing.assert_close(
                            v,
                            reloaded_state[k],
                            msg=lambda x: f"DabDetrModel: Tensor {k}: {x}. Key {k} was serialized: {k in serialized_keys}. If `False`, this means it was probably aliased and safetensors removed it. If `True` it means `_init_weights` overwrote that key",
                        )

            # Checking there was no complain of missing weights
            self.assertEqual(
                infos["missing_keys"],
                set(),
                "Given that the loaded weights are the same, the issue is in `tie_weights`: it tied these keys and removed them from serialization. But because of tiying (hardcoded or not) the previous check is fine.\
                    This can happen if `save_pretrained` remove the targets and not the keys from serialiazation, or you hardcoded `self.xxx = yyy` thus forcing to always tie -> they are removed from serialization.",
            )

    # TODO: check if this works again for PyTorch 2.x.y
    @unittest.skip(reason="Got `CUDA error: misaligned address` with PyTorch 2.0.0.")