fix return_result bug

c77d0dee · Steffy-zxf · 8666b7b8 · c77d0dee
隐藏空白更改
内联并排

Showing with 1 addition and 1 deletion

paddlehub/finetune/task/classifier_task.py paddlehub/finetune/task/classifier_task.py +1 -1

未找到文件。
--- a/paddlehub/finetune/task/classifier_task.py
+++ b/paddlehub/finetune/task/classifier_task.py
@@ -150,7 +150,7 @@ class ClassifierTask(BaseTask):
        results = []
        for batch_state in run_states:
            batch_result = batch_state.run_results
-            batch_infer = np.argmax(batch_result, axis=2)[0]
+            batch_infer = np.argmax(batch_result[0], axis=1)
            results += [id2label[sample_infer] for sample_infer in batch_infer]
        return results