def argmax_batch(vecs: nd.NDArray):
# _, idx = torch.max(vecs, 1)
# return idx
return vecs.argmax(axis=1)
def argmax(vec: nd.NDArray):
# _, idx = torch.max(vec, 1)
# return to_scalar(idx)
return vec.argmax(1)
def argmax_batch(vecs: nd.NDArray):
# _, idx = torch.max(vecs, 1)
# return idx
return vecs.argmax(axis=1)
def argmax(vec: nd.NDArray):
# _, idx = torch.max(vec, 1)
# return to_scalar(idx)
return vec.argmax(1)