fix(encoders): fix the broken transformers

jina-ai · May 11, 2020 · 5b62c31 · 5b62c31
1 parent 80224f7
commit 5b62c31
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jina/executors/encoders/nlp/transformer.py b/jina/executors/encoders/nlp/transformer.py
@@ -65,7 +65,7 @@ def encode(self, data: 'np.ndarray', *args, **kwargs) -> 'np.ndarray':
             _mask_ids_batch = self.tensor2array(mask_ids_batch)
             _seq_output = self.tensor2array(seq_output)
             if self.pooling_strategy == 'cls':
-                if hasattr(self._tokenizer, 'cls_token'):
+                if hasattr(self._tokenizer, 'cls_token') and len(extra_output) > 0:
                     output = self.tensor2array(extra_output[0])
                 else:
                     output = reduce_cls(_seq_output, _mask_ids_batch, self.cls_pos)