Update finetuned_model_to_module.md
@@ -192,7 +192,7 @@ def predict(self, data, return_result=False, accelerate_mode=True):
prediction = []
for batch_result in results:
# get predict index
- batch_result = np.argmax(batch_result, axis=2)[0]
+ batch_result = np.argmax(batch_result[0], axis=1)
batch_result = batch_result.tolist()
prediction += batch_result
return prediction
