Sanitize export outputs for vits

src/otx/algo/classification/vit.py

@@ -94,7 +94,6 @@ def head_forward_fn(self, x: torch.Tensor) -> torch.Tensor:
         x = self.model.backbone.norm(x)
         if self.model.neck is not None:
             x = self.model.neck(x)
-
         # Head
         cls_token = x[:, 0]
         layer_output = [None, cls_token]
@@ -140,14 +139,18 @@ def _forward_explain_image_classifier(
             scores = pred_results.unbind(0)
             labels = logits.argmax(-1, keepdim=True).unbind(0)
 
-        return {
+        outputs =  {
             "logits": logits,
             "feature_vector": feature_vector,
             "saliency_map": saliency_map,
-            "scores": scores,
-            "labels": labels,
         }
 
+        if not torch._C._is_tracing():
+            outputs["scores"] = scores
+            outputs["labels"] = labels
+
+        return outputs
+
     def get_explain_fn(self) -> Callable:
         """Returns explain function."""
         explainer = ViTReciproCAM(