out = model(input_ids.as_in_context(ctx),
type_ids.as_in_context(ctx),
valid_length.astype("float32").as_in_context(ctx))
values, indices = mx.nd.topk(out, k=1, ret_typ="both")
value_list.extend(values.asnumpy().reshape(-1).tolist())
index_list.extend(indices.asnumpy().reshape(-1).tolist())
After Change
logging.info("Time cost=%.2fs, throughput=%.2f samples/s", toc - tic,
dev_batch_size * len(loader_test) / (toc - tic))
// write result to a file.
filename = args.task_name + segment.replace("test", "") + ".csv"
test_path = os.path.join(args.output_dir, filename)
with io.open(test_path, "w", encoding="utf-8") as f:
f.write(u"index\tprediction\n")
for i, pred in enumerate(results):
