a = torch.rand(2000,3000).cuda()
time_start = time.time()
for i in range(2048):
max_index = torch.argmax(a)
print(“run %.2f” % (time.time() - time_start))
the top 1024 index in for loop return result quickly， but the latter is slower，
how to fix this problem?