DataLoader type inference tests

FluxML · Jun 12, 2020 · a887426 · a887426
1 parent c9219fb
commit a887426
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/src/data/dataloader.jl b/src/data/dataloader.jl
@@ -105,6 +105,6 @@ function _nobs(data::Union{Tuple, NamedTuple})
 end
 
 _getobs(data::AbstractArray, i) = selectdim(data, ndims(data), i)
-_getobs(data::Union{Tuple, NamedTuple}, i) = map(x -> _getobs(x, i), data)
+_getobs(data::Union{Tuple, NamedTuple}, i) = map(Base.Fix2(_getobs, i), data)
 
-Base.eltype(d::DataLoader{D}) where D = D
+Base.eltype(::DataLoader{D}) where D = D
diff --git a/test/data.jl b/test/data.jl
@@ -3,6 +3,7 @@
     Y = [1:5;]
 
     d = DataLoader(X, batchsize=2)
+    @inferred first(d)
     batches = collect(d)
     @test eltype(batches) == eltype(d) == typeof(X)
     @test length(batches) == 3
@@ -11,20 +12,23 @@
     @test batches[3] == X[:,5:5]
 
     d = DataLoader(X, batchsize=2, partial=false)
+    @inferred first(d)
     batches = collect(d)
     @test eltype(batches) == eltype(d) == typeof(X)
     @test length(batches) == 2
     @test batches[1] == X[:,1:2]
     @test batches[2] == X[:,3:4]
 
     d = DataLoader((X,), batchsize=2, partial=false)
+    @inferred first(d)
     batches = collect(d)
     @test eltype(batches) == eltype(d) == Tuple{typeof(X)}
     @test length(batches) == 2
     @test batches[1] == (X[:,1:2],)
     @test batches[2] == (X[:,3:4],)
 
     d = DataLoader((X, Y), batchsize=2)
+    @inferred first(d)
     batches = collect(d)
     @test eltype(batches) == eltype(d) == Tuple{typeof(X), typeof(Y)}
     @test length(batches) == 3
@@ -40,6 +44,7 @@
 
     # test with NamedTuple
     d = DataLoader((x=X, y=Y), batchsize=2)
+    @inferred first(d)
     batches = collect(d)
     @test eltype(batches) == eltype(d) == NamedTuple{(:x, :y), Tuple{typeof(X), typeof(Y)}}
     @test length(batches) == 3