fix partition

Author: IanButterworth

src/impl/conv_im2col.jl

@@ -45,7 +45,7 @@ function conv_im2col!(
     N = channels_out(cdims)
     K = prod(kernel_size(cdims))*channels_in(cdims)
 
-    parts = collect(Iterators.partition(axes(x, 5), ntasks))
+    parts = collect(Iterators.partition(axes(x, 5), ceil(Int, size(x, 5) / ntasks)))
 
     @sync for task_n in 1:ntasks
         Threads.@spawn begin
@@ -150,7 +150,7 @@ function ∇conv_data_im2col!(
     N = prod(kernel_size(cdims))*channels_in(cdims)
     K = channels_out(cdims)
 
-    parts = collect(Iterators.partition(axes(dx, 5), ntasks))
+    parts = collect(Iterators.partition(axes(dx, 5), ceil(Int, size(dx, 5) / ntasks)))
 
     @sync for task_n in 1:ntasks
         Threads.@spawn begin

src/impl/depthwiseconv_im2col.jl

@@ -26,7 +26,7 @@ function depthwiseconv_im2col!(
     N = channel_multiplier(cdims)
     K = prod(kernel_size(cdims))
 
-    parts = collect(Iterators.partition(axes(y)[end], ntasks))
+    parts = collect(Iterators.partition(axes(y)[end], ceil(Int, size(y, 5) / ntasks)))
 
     dcdims = DenseConvDims(cdims)
 
@@ -114,7 +114,7 @@ function ∇depthwiseconv_data_im2col!(
     N = prod(kernel_size(cdims))
     K = channel_multiplier(cdims)
 
-    parts = collect(Iterators.partition(axes(dx)[end], ntasks))
+    parts = collect(Iterators.partition(axes(dx)[end], ceil(Int, size(dx, 5) / ntasks)))
 
     @sync for task_n in 1:ntasks
         Threads.@spawn begin