Softmax fix

[rw404]

• Softmax layer in eval mode added;
• model.train(), model.eval() in train.py placed for better performance

Fixes from issue #51

[twofeetcat]

hello，I can probably understand what you mean and follow your code to repeat it, but when calculating the loss, I wonder whether the following error occurred because my real label had 20 classes（the value of y_batch is 0, 1, ..., 19, shape of y_batch is [batch_size, 256, 256]） Error reported as follows: Traceback (most recent call last): File "train.py", line 141, in <module> loss.backward() File "D:\software\anaconda\lib\site-packages\torch\_tensor.py", line 489, in backward self, gradient, retain_graph, create_graph, inputs=inputs File "D:\software\anaconda\lib\site-packages\torch\autograd\__init__.py", line 199, in backward allow_unreachable=True, accumulate_grad=True) # Calls into the C++ engine to run the backward pass RuntimeError: Unable to find a valid cuDNN algorithm to run convolution
But when I changed all values greater than 0 in y_batch to 1, that is, 20 classes to 0 and 1, the program did not have this error