From 98f6e1ee87e83355affd7951d7d69347e1945294 Mon Sep 17 00:00:00 2001 From: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> Date: Thu, 26 May 2022 09:57:12 -0400 Subject: [PATCH] Fix model parallelism test (#17439) --- tests/test_modeling_common.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 9a7777af4a..4b9227cbfd 100755 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -2203,7 +2203,7 @@ class ModelTesterMixin: @require_torch_gpu def test_cpu_offload(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - if config.num_hidden_layers < 5: + if isinstance(getattr(config, "num_hidden_layers", None), int) and config.num_hidden_layers < 5: config.num_hidden_layers = 5 for model_class in self.all_model_classes: @@ -2236,7 +2236,7 @@ class ModelTesterMixin: @require_torch_multi_gpu def test_model_parallelism(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - if config.num_hidden_layers < 5: + if isinstance(getattr(config, "num_hidden_layers", None), int) and config.num_hidden_layers < 5: config.num_hidden_layers = 5 for model_class in self.all_model_classes: