From 3b0c1488af95f2b9d05264bee4a3f4dbedb1c1d8 Mon Sep 17 00:00:00 2001
From: "Kang, Harim" <harim.kang@intel.com>
Date: Mon, 20 Mar 2023 14:21:36 +0900
Subject: [PATCH] Fix VisionTransformeroutput check

---
 otx/mpa/cls/stage.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/otx/mpa/cls/stage.py b/otx/mpa/cls/stage.py
index 08d5d051eb8..eeefaa14db0 100644
--- 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