fix when there's few cols

Author: IanButterworth

src/impl/conv_im2col.jl

@@ -47,7 +47,7 @@ function conv_im2col!(
 
     parts = collect(Iterators.partition(axes(x, 5), ceil(Int, size(x, 5) / ntasks)))
 
-    @sync for task_n in 1:ntasks
+    @sync for task_n in eachindex(parts)
         Threads.@spawn begin
             col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
             for batch_idx in parts[task_n]
@@ -152,7 +152,7 @@ function ∇conv_data_im2col!(
 
     parts = collect(Iterators.partition(axes(dx, 5), ceil(Int, size(dx, 5) / ntasks)))
 
-    @sync for task_n in 1:ntasks
+    @sync for task_n in eachindex(parts)
         Threads.@spawn begin
             col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
             for batch_idx in parts[task_n]

src/impl/depthwiseconv_im2col.jl

@@ -30,7 +30,7 @@ function depthwiseconv_im2col!(
 
     dcdims = DenseConvDims(cdims)
 
-    @sync for task_n in 1:ntasks
+    @sync for task_n in eachindex(parts)
         Threads.@spawn begin
             col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
             for batch_idx in parts[task_n]
@@ -117,7 +117,7 @@ function ∇depthwiseconv_data_im2col!(
 
     parts = collect(Iterators.partition(axes(dx)[end], ceil(Int, size(dx, 5) / ntasks)))
 
-    @sync for task_n in 1:ntasks
+    @sync for task_n in eachindex(parts)
         Threads.@spawn begin
             col_slice = col_slice = view(col, :, :, task_n) # col_slice is a task-local workspace
             for batch_idx in parts[task_n]