fix benchmark non standard model (huggingface#5801)

zsc · Jul 16, 2020 · aefc0c0 · aefc0c0
1 parent 8ce610b
commit aefc0c0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/transformers/benchmark/benchmark.py b/src/transformers/benchmark/benchmark.py
@@ -88,7 +88,7 @@ def _prepare_inference_func(self, model_name: str, batch_size: int, sequence_len
         if self.args.torchscript:
             config.torchscript = True
 
-        has_model_class_in_config = hasattr(config, "architecture") and len(config.architectures) > 1
+        has_model_class_in_config = hasattr(config, "architectures") and len(config.architectures) > 0
         if not self.args.only_pretrain_model and has_model_class_in_config:
             try:
                 model_class = config.architectures[0]
@@ -138,7 +138,7 @@ def encoder_forward():
     def _prepare_train_func(self, model_name: str, batch_size: int, sequence_length: int) -> Callable[[], None]:
         config = self.config_dict[model_name]
 
-        has_model_class_in_config = hasattr(config, "architecture") and len(config.architectures) > 1
+        has_model_class_in_config = hasattr(config, "architectures") and len(config.architectures) > 0
         if not self.args.only_pretrain_model and has_model_class_in_config:
             try:
                 model_class = config.architectures[0]

diff --git a/src/transformers/benchmark/benchmark_tf.py b/src/transformers/benchmark/benchmark_tf.py
@@ -132,7 +132,7 @@ def _prepare_inference_func(self, model_name: str, batch_size: int, sequence_len
         if self.args.fp16:
             raise NotImplementedError("Mixed precision is currently not supported.")
 
-        has_model_class_in_config = hasattr(config, "architecture") and len(config.architectures) > 1
+        has_model_class_in_config = hasattr(config, "architectures") and len(config.architectures) > 0
         if not self.args.only_pretrain_model and has_model_class_in_config:
             try:
                 model_class = "TF" + config.architectures[0]  # prepend 'TF' for tensorflow model
@@ -172,7 +172,7 @@ def _prepare_train_func(self, model_name: str, batch_size: int, sequence_length:
         if self.args.fp16:
             raise NotImplementedError("Mixed precision is currently not supported.")
 
-        has_model_class_in_config = hasattr(config, "architecture") and len(config.architectures) > 1
+        has_model_class_in_config = hasattr(config, "architectures") and len(config.architectures) > 0
         if not self.args.only_pretrain_model and has_model_class_in_config:
             try:
                 model_class = "TF" + config.architectures[0]  # prepend 'TF' for tensorflow model