From b210c83a78022226ce48402cd67d8c8da7afbd8d Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Fri, 13 Jan 2023 10:03:19 +0100 Subject: [PATCH] Fix `torchscript` tests for `AltCLIP` (#21102) fix torchscript tests for AltCLIP Co-authored-by: ydshieh --- tests/models/altclip/test_modeling_altclip.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/models/altclip/test_modeling_altclip.py b/tests/models/altclip/test_modeling_altclip.py index 63f3d62171..c22b9a4684 100755 --- a/tests/models/altclip/test_modeling_altclip.py +++ b/tests/models/altclip/test_modeling_altclip.py @@ -490,6 +490,15 @@ class AltCLIPModelTest(ModelTesterMixin, unittest.TestCase): model_state_dict = model.state_dict() loaded_model_state_dict = loaded_model.state_dict() + non_persistent_buffers = {} + for key in loaded_model_state_dict.keys(): + if key not in model_state_dict.keys(): + non_persistent_buffers[key] = loaded_model_state_dict[key] + + loaded_model_state_dict = { + key: value for key, value in loaded_model_state_dict.items() if key not in non_persistent_buffers + } + self.assertEqual(set(model_state_dict.keys()), set(loaded_model_state_dict.keys())) models_equal = True