Fix VisionTransformeroutput check

openvinotoolkit · Mar 20, 2023 · 3b0c148 · 3b0c148
1 parent 85f033a
commit 3b0c148
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/otx/mpa/cls/stage.py b/otx/mpa/cls/stage.py
@@ -92,7 +92,7 @@ def configure_in_channel(cfg):
         if isinstance(output, (tuple, list)):
             output = output[-1]
 
-        if layer.__class__.__name__ in TRANSFORMER_BACKBONES:
+        if layer.__class__.__name__ in TRANSFORMER_BACKBONES and isinstance(output, (tuple, list)):
             # mmcls.VisionTransformer outputs Tuple[List[...]] and the last index of List is the final logit.
             _, output = output