Fix image painting

[rtyasdf]

There is an attempt to save images in train.py and test.py, and main step towards making such an image is separation output of the model by threshold 0.5 . Such action make sense if model returns probabilities, but final layer in medt_net and ResAxialAttentionUNet classes is Conv2d, which returns arbitrary floats, not bounded by [0; 1]. So I added softmax to directly convert output to probabilities when going through validation and test datasets.

Loss calculation by the way is right, because torch.nn.functional.cross_entropy expect to see logits, not probabilites. Directly thresholding logits just can make one wondering why loss decreasing but images still looks bad.