feat(transformers): debug video classify pipe

wcrzlh · Jan 22, 2025 · 82602e3 · 82602e3
1 parent d8895dd
commit 82602e3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mindone/transformers/pipelines/video_classification.py b/mindone/transformers/pipelines/video_classification.py
@@ -154,7 +154,7 @@ def postprocess(self, model_outputs, top_k=5, function_to_apply="softmax"):
         if top_k > self.model.config.num_labels:
             top_k = self.model.config.num_labels
 
-        if self.framework == "pt":
+        if self.framework == "ms":
             if function_to_apply == "softmax":
                 probs = model_outputs.logits[0].softmax(-1)
             elif function_to_apply == "sigmoid":