Its good that you removed the softmax, so now you have only 1 softmax (in CrossEntropyLoss), as expected.
Adding another Relu might help, but its not a game changer because nonlinearity already exists in your model.
Does the nans from your loss? or from your accuracy metrics?
Since its a classification challenge, i would say that CategoricalAccuracy is needed
Ive taken this implementation from lpd 1
    def categorical_correct(y_pred, y_true):
        indices = T.max(y_pred, 1)[1]
        correct = T.eq(indices, y_true).view(-1)
        return correct.float().sum()

so in your example, you should do:
train_correct += categorical_correct(predicted, labels)
and for val:
correct_val += categorical_correct(predicted, labels)
Roy