using nn.Module(X).argmax(1) - get IndexError

Hello there, I'm student of NN course, I'm try to implement FFNN (or TDNN) to work on prediction of AR(2)-model, im using PyTorch example, and on my data and NN architecture i got pred.argmax(1) - error:

Traceback (most recent call last):
  File "/home/b0r1ngx/PycharmProjects/ArtificialNeuroNets/group_00201/lab01/lab01_pytorch.py", line 116, in <module>
    first_method()
  File "/home/b0r1ngx/PycharmProjects/ArtificialNeuroNets/group_00201/lab01/lab01_pytorch.py", line 87, in first_method
    test(test_data, time_delay_nn, loss_function)
  File "/home/b0r1ngx/PycharmProjects/ArtificialNeuroNets/group_00201/lab01/lab01_pytorch.py", line 70, in test
    c1 = pred.argmax(1) == y_pred
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

where it's used in you're examples:
here in test_loop function - https://github.com/pytorch/tutorials/blob/master/beginner_source/basics/optimization_tutorial.py

I'm also doesn't think that i get best Hyperparameters / loss_function / optimizer - cos i get bad Accuracy / Avg loss in my case,
please help me with that:
U can check my code here:
(now im using how its recommended -1 or 0, but there is always 0)
https://github.com/b0r1ngx/ArtificialNeuroNets/blob/master/group_00201/lab01/lab01_pytorch.py

Thanks!

cc @jerryzh168 @z-a-f @vkuzo

Start with group_00201/lab01/lab01_pytorch.py, especially the test function around line 70 and first_method around line 87, then compare the prediction handling with the linked PyTorch optimization tutorial's test_loop. Check the prediction tensor shape and the reported accuracy and average loss; done means the IndexError and the unexpected metrics are explained for this example.