test.py has a high time cost.

[StuLiu]

line 142 in test.py: "pred = pred.cpu().numpy().squeeze().argmax(0)" should change to "pred = pred.squeeze().argmax(0).cpu().numpy()".