revert selectdim

selectdim can lead to type instability, see https://discourse.julialang.org/t/why-selectdim-is-type-instable/25271/5
FluxML · Jun 16, 2020 · 9078f85 · 9078f85
1 parent 1dbaf32
commit 9078f85
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/data/dataloader.jl b/src/data/dataloader.jl
@@ -104,7 +104,7 @@ function _nobs(data::Union{Tuple, NamedTuple})
     return n
 end
 
-_getobs(data::AbstractArray, i) = selectdim(data, ndims(data), i)
+_getobs(data::AbstractArray, i) = data[ntuple(i -> Colon(), Val(ndims(data) - 1))..., i]
 _getobs(data::Union{Tuple, NamedTuple}, i) = map(Base.Fix2(_getobs, i), data)
 
 Base.eltype(::DataLoader{D}) where D = D