argmax won’t solve this issue, which points to a mismatch between the valid target indices and the model output shape.
Assuming you are using nn.CrossEntropyLoss the output should have the shape [batch_size, nb_classes] while the target should have the shape [batch_size] and contain class indices in the range [0, nb_classes-1]. Since your target contains values as large as 156 you would be dealing with at least 157 classes (valid indices are [0, 156]) and the model output should thus have the shape [batch_size, 157].
While this doesn’t fit your description of working on a regression task, the error would fit into an indexing issue in nn.CrossEntropyLoss.