test.py error

[gu-yan]

An error occurs with code below: test.py
pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=1)
decoded = loader.decode_segmap(pred[0]).

I use python3.5, pytorch 0.2.0 and numpy 1.13.1.
It works with changes below:
pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0)
decoded = loader.decode_segmap(pred)

[jetxa]

try this code

pred = outputs.data.max(1)[1].cpu().numpy()

[gu-yan]

@zhxt0 thank you, it works