Fix _segmentationloss for 3D images

[itan1]

The _segmentationloss reshapes the array before passing it to the logitcrossentropy loss from Flux. During reshaping it is assumed that ys is 4 dimensional. This won't work for 3D segmentation tasks, where ys will be 5 dimensional.

Previously:

julia> _segmentationloss(zeros(10, 10, 10, 3, 5), zeros(10, 10, 10, 3, 5))
ERROR: DimensionMismatch: loss function expects size(ŷ) = (1000, 3, 5) to match size(y) = (500, 10, 3)
Stacktrace:
 [1] _check_sizes(ŷ::Array{Float64, 3}, y::Array{Float64, 3})
   @ Flux.Losses ~/.julia/packages/Flux/KkC79/src/losses/utils.jl:31
 [2] logitcrossentropy(ŷ::Array{Float64, 3}, y::Array{Float64, 3}; dims::Int64, agg::typeof(Statistics.mean))
   @ Flux.Losses ~/.julia/packages/Flux/KkC79/src/losses/functions.jl:254
 [3] _segmentationloss(ypreds::Array{Float64, 5}, ys::Array{Float64, 5}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
   @ Main ./REPL[245]:5
 [4] _segmentationloss(ypreds::Array{Float64, 5}, ys::Array{Float64, 5})
   @ Main ./REPL[245]:1
 [5] top-level scope
   @ REPL[258]:1
 [6] top-level scope
   @ ~/.julia/packages/CUDA/DfvRa/src/initialization.jl:52

Now

julia> _segmentationloss(zeros(10, 10, 10, 3, 5), zeros(10, 10, 10, 3, 5))
0.0

PR Checklist

• [x] Tests are added
• [x] Documentation, if applicable -> no update needed

[lorenzoh]

Good catch and thanks for the fix!