diff --git a/Project.toml b/Project.toml
index 0fa3946..4a23524 100644
--- a/Project.toml
+++ b/Project.toml
@@ -33,6 +33,7 @@ julia = "1.9"
 [extras]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
@@ -42,4 +43,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [targets]
-test = ["CUDA", "cuDNN", "LinearAlgebra", "MLJBase", "Random", "StableRNGs", "StatisticalMeasures", "StatsBase", "Test"]
+test = ["CUDA", "cuDNN", "LinearAlgebra", "Logging", "MLJBase", "Random", "StableRNGs", "StatisticalMeasures", "StatsBase", "Test"]
diff --git a/src/core.jl b/src/core.jl
index b3df605..f866ee5 100644
--- a/src/core.jl
+++ b/src/core.jl
@@ -276,15 +276,22 @@ input `X` and target `y` in the form required by
 by `model.batch_size`.)
 
 """
-function collate(model, X, y)
+function collate(model, X, y, verbosity)
     row_batches = Base.Iterators.partition(1:nrows(y), model.batch_size)
-    Xmatrix = reformat(X)
+    Xmatrix = _f32(reformat(X), verbosity)
     ymatrix = reformat(y)
     return [_get(Xmatrix, b) for b in row_batches], [_get(ymatrix, b) for b in row_batches]
 end
-function collate(model::NeuralNetworkBinaryClassifier, X, y)
+function collate(model::NeuralNetworkBinaryClassifier, X, y, verbosity)
     row_batches = Base.Iterators.partition(1:nrows(y), model.batch_size)
-    Xmatrix = reformat(X)
+    Xmatrix = _f32(reformat(X), verbosity)
     yvec = (y .== classes(y)[2])' # convert to boolean
     return [_get(Xmatrix, b) for b in row_batches], [_get(yvec, b) for b in row_batches]
 end
+
+_f32(x::AbstractArray{Float32}, verbosity) = x
+function _f32(x::AbstractArray, verbosity)
+    verbosity > 0 && @info "MLJFlux: converting input data to Float32"
+    return Float32.(x)
+end
+
diff --git a/src/encoders.jl b/src/encoders.jl
index de51714..f396142 100644
--- a/src/encoders.jl
+++ b/src/encoders.jl
@@ -19,11 +19,11 @@ function ordinal_encoder_fit(X; featinds)
         feat_col = Tables.getcolumn(Tables.columns(X), i)
         feat_levels = levels(feat_col)
         # Check if feat levels is already ordinal encoded in which case we skip
-        (Set([float(i) for i in 1:length(feat_levels)]) == Set(feat_levels)) && continue
+        (Set([Float32(i) for i in 1:length(feat_levels)]) == Set(feat_levels)) && continue
         # Compute the dict using the given feature_mapper function
         mapping_matrix[i] =
-            Dict{Any, AbstractFloat}(
-                value => float(index) for (index, value) in enumerate(feat_levels)
+            Dict{eltype(feat_levels), Float32}(
+                value => Float32(index) for (index, value) in enumerate(feat_levels)
             )
     end
     return mapping_matrix
@@ -67,7 +67,7 @@ function ordinal_encoder_transform(X, mapping_matrix)
             test_levels = levels(col)
             check_unkown_levels(train_levels, test_levels)
             level2scalar = mapping_matrix[ind]
-            new_col = recode(col, level2scalar...)
+            new_col = recode(unwrap.(col), level2scalar...)
             push!(new_feats, new_col)
         else
             push!(new_feats, col)
diff --git a/src/entity_embedding.jl b/src/entity_embedding.jl
index 3e16ea2..313e3e6 100644
--- a/src/entity_embedding.jl
+++ b/src/entity_embedding.jl
@@ -36,7 +36,6 @@ julia> output = embedder(batch)
 ```
 """ # 1. Define layer struct to hold parameters
 struct EntityEmbedder{A1 <: AbstractVector, A2 <: AbstractVector, I <: Integer}
-    
     embedders::A1
     modifiers::A2       # applied on the input before passing it to the embedder
     numfeats::I
@@ -44,7 +43,7 @@ end
 
 # 2. Define the forward pass (i.e., calling an instance of the layer)
 (m::EntityEmbedder)(x) =
-    vcat([m.embedders[i](m.modifiers[i](x, i)) for i in 1:m.numfeats]...)
+    (vcat([m.embedders[i](m.modifiers[i](x, i)) for i in 1:m.numfeats]...))
 
 # 3. Define the constructor which initializes the parameters and returns the instance
 function EntityEmbedder(entityprops, numfeats; init = Flux.randn32)
diff --git a/src/mlj_model_interface.jl b/src/mlj_model_interface.jl
index e7d840f..f3d2c9c 100644
--- a/src/mlj_model_interface.jl
+++ b/src/mlj_model_interface.jl
@@ -66,35 +66,35 @@ function MLJModelInterface.fit(model::MLJFluxModel,
     X,
     y)
     # GPU and rng related variables
-    move = Mover(model.acceleration)
+    move = MLJFlux.Mover(model.acceleration)
     rng = true_rng(model)
 
     # Get input properties
     shape = MLJFlux.shape(model, X, y)
-    cat_inds = get_cat_inds(X)
+    cat_inds = MLJFlux.get_cat_inds(X)
     pure_continuous_input = isempty(cat_inds)
 
     # Decide whether to enable entity embeddings (e.g., ImageClassifier won't)
-    enable_entity_embs = is_embedding_enabled(model) && !pure_continuous_input
+    enable_entity_embs = MLJFlux.is_embedding_enabled(model) && !pure_continuous_input
 
     # Prepare entity embeddings inputs and encode X if entity embeddings enabled
     featnames = []
     if enable_entity_embs
-        X = convert_to_table(X)
+        X = MLJFlux.convert_to_table(X)
         featnames = Tables.schema(X).names
     end
 
-    # entityprops is (index = cat_inds[i], levels = num_levels[i], newdim = newdims[i]) 
+    # entityprops is (index = cat_inds[i], levels = num_levels[i], newdim = newdims[i])
     # for each categorical feature
     default_embedding_dims = enable_entity_embs ? model.embedding_dims : Dict{Symbol, Real}()
     entityprops, entityemb_output_dim =
-        prepare_entityembs(X, featnames, cat_inds, default_embedding_dims)
-    X, ordinal_mappings = ordinal_encoder_fit_transform(X; featinds = cat_inds)
+        MLJFlux.prepare_entityembs(X, featnames, cat_inds, default_embedding_dims)
+    X, ordinal_mappings = MLJFlux.ordinal_encoder_fit_transform(X; featinds = cat_inds)
 
     ## Construct model chain
     chain =
         (!enable_entity_embs) ? construct_model_chain(model, rng, shape, move) :
-        construct_model_chain_with_entityembs(
+        MLJFlux.construct_model_chain_with_entityembs(
             model,
             rng,
             shape,
@@ -103,8 +103,8 @@ function MLJModelInterface.fit(model::MLJFluxModel,
             entityemb_output_dim,
         )
 
-    # Format data as needed by Flux and move to GPU 
-    data = move.(collate(model, X, y))
+    # Format data as needed by Flux and move to GPU
+    data = move.(MLJFlux.collate(model, X, y, verbosity))
 
     # Test chain works (as it may be custom)
     x = data[1][1]
@@ -143,7 +143,7 @@ function MLJModelInterface.fit(model::MLJFluxModel,
         featnames,
     )
 
-    # Prepare fitresult 
+    # Prepare fitresult
     fitresult =
         MLJFlux.fitresult(model, Flux.cpu(chain), y, ordinal_mappings, embedding_matrices)
 
@@ -216,7 +216,7 @@ function MLJModelInterface.update(model::MLJFluxModel,
             chain = construct_model_chain(model, rng, shape, move)
         end
         # reset `optimiser_state`:
-        data = move.(collate(model, X, y))
+        data = move.(collate(model, X, y, verbosity))
         regularized_optimiser, optimiser_state =
             prepare_optimiser(data, model, chain)
         epochs = model.epochs
diff --git a/test/classifier.jl b/test/classifier.jl
index 5d7f144..2dbafef 100644
--- a/test/classifier.jl
+++ b/test/classifier.jl
@@ -4,18 +4,17 @@ seed!(1234)
 N = 300
 Xm = MLJBase.table(randn(Float32, N, 5));   # purely numeric
 X = (; Tables.columntable(Xm)...,
-    Column1 = repeat([1.0, 2.0, 3.0, 4.0, 5.0], Int(N / 5)),
+    Column1 = repeat(Float32[1.0, 2.0, 3.0, 4.0, 5.0], Int(N / 5)),
     Column2 = categorical(repeat(['a', 'b', 'c', 'd', 'e'], Int(N / 5))),
     Column3 = categorical(repeat(["b", "c", "d", "f", "f"], Int(N / 5)), ordered = true),
-    Column4 = repeat([1.0, 2.0, 3.0, 4.0, 5.0], Int(N / 5)),
-    Column5 = randn(N),
+    Column4 = repeat(Float32[1.0, 2.0, 3.0, 4.0, 5.0], Int(N / 5)),
+    Column5 = randn(Float32, N),
     Column6 = categorical(
         repeat(["group1", "group1", "group2", "group2", "group3"], Int(N / 5)),
     ),
 )
 
-
-ycont = 2 * X.x1 - X.x3 + 0.1 * rand(N)
+ycont = 2 * X.x1 - X.x3 + 0.1 * rand(Float32, N)
 m, M = minimum(ycont), maximum(ycont)
 _, a, b, _ = range(m, stop = M, length = 4) |> collect
 y = map(ycont) do η
@@ -111,7 +110,8 @@ end
 
 # check different resources (CPU1, CUDALibs, etc)) give about the same loss:
 reference = losses[1]
-@test all(x -> abs(x - reference) / reference < 1e-4, losses[2:end])
+println("losses for each resource: $losses")
+@test all(x -> abs(x - reference) / reference < 0.03, losses[2:end])
 
 
 # # NEURAL NETWORK BINARY CLASSIFIER
@@ -126,7 +126,7 @@ end
 seed!(1234)
 N = 300
 X = MLJBase.table(rand(Float32, N, 4));
-ycont = 2 * X.x1 - X.x3 + 0.1 * rand(N)
+ycont = Float32.(2 * X.x1 - X.x3 + 0.1 * rand(N))
 m, M = minimum(ycont), maximum(ycont)
 _, a, _ = range(m, stop = M, length = 3) |> collect
 y = map(ycont) do η
diff --git a/test/core.jl b/test/core.jl
index 4bfc400..94eb1e1 100644
--- a/test/core.jl
+++ b/test/core.jl
@@ -14,23 +14,28 @@ rowvec(y::Vector) = reshape(y, 1, length(y))
 end
 
 @testset "collate" begin
-    # NeuralNetworRegressor:
-    Xmatrix = broadcast(x->round(x, sigdigits=2), rand(stable_rng, 10, 3))
+    Xmatrix = broadcast(x->round(x, sigdigits=2), rand(stable_rng, Float32, 10, 3))
+    Xmat_f64 = Float64.(Xmatrix)
     # convert to a column table:
     X = MLJBase.table(Xmatrix)
+    X_64 = MLJBase.table(Xmat_f64)
 
+    # NeuralNetworRegressor:
     y = rand(stable_rng, Float32, 10)
     model = MLJFlux.NeuralNetworkRegressor()
     model.batch_size= 3
-    @test MLJFlux.collate(model, X, y) ==
+    @test MLJFlux.collate(model, X, y, 1) == MLJFlux.collate(model, X_64, y, 1) ==
         ([Xmatrix'[:,1:3], Xmatrix'[:,4:6], Xmatrix'[:,7:9], Xmatrix'[:,10:10]],
          rowvec.([y[1:3], y[4:6], y[7:9], y[10:10]]))
+    @test_logs (:info,) MLJFlux.collate(model, X_64, y, 1)
+    @test_logs min_level=Logging.Info MLJFlux.collate(model, X, y, 1)
+    @test_logs min_level=Logging.Info MLJFlux.collate(model, X, y, 0)
 
     # NeuralNetworClassifier:
     y = categorical(['a', 'b', 'a', 'a', 'b', 'a', 'a', 'a', 'b', 'a'])
     model = MLJFlux.NeuralNetworkClassifier()
     model.batch_size = 3
-    data = MLJFlux.collate(model, X, y)
+    data = MLJFlux.collate(model, X, y, 1)
 
     @test data == ([Xmatrix'[:,1:3], Xmatrix'[:,4:6],
                     Xmatrix'[:,7:9], Xmatrix'[:,10:10]],
@@ -42,13 +47,13 @@ end
     y = MLJBase.table(ymatrix) # a rowaccess table
     model = MLJFlux.NeuralNetworkRegressor()
     model.batch_size= 3
-    @test MLJFlux.collate(model, X, y) ==
+    @test MLJFlux.collate(model, X, y, 1) ==
         ([Xmatrix'[:,1:3], Xmatrix'[:,4:6], Xmatrix'[:,7:9], Xmatrix'[:,10:10]],
          rowvec.([ymatrix'[:,1:3], ymatrix'[:,4:6], ymatrix'[:,7:9],
                   ymatrix'[:,10:10]]))
 
     y = Tables.columntable(y) # try a columnaccess table
-    @test MLJFlux.collate(model, X, y) ==
+    @test MLJFlux.collate(model, X, y, 1) ==
         ([Xmatrix'[:,1:3], Xmatrix'[:,4:6], Xmatrix'[:,7:9], Xmatrix'[:,10:10]],
          rowvec.([ymatrix'[:,1:3], ymatrix'[:,4:6],
                   ymatrix'[:,7:9], ymatrix'[:,10:10]]))
@@ -58,7 +63,7 @@ end
     y = categorical(['a', 'b', 'a', 'a', 'b', 'a', 'a', 'a', 'b', 'a'])
     model = MLJFlux.ImageClassifier(batch_size=2)
 
-    data = MLJFlux.collate(model, Xmatrix, y)
+    data = MLJFlux.collate(model, Xmatrix, y, 1)
     @test  first.(data) == (Float32.(cat(Xmatrix[1], Xmatrix[2], dims=4)),
                             rowvec.([1 0;0 1]))
 
diff --git a/test/encoders.jl b/test/encoders.jl
index 7bd5da9..50b93b6 100644
--- a/test/encoders.jl
+++ b/test/encoders.jl
@@ -12,8 +12,8 @@
     @test map[2] == Dict('a' => 1, 'b' => 2, 'c' => 3, 'd' => 4, 'e' => 5)
     @test map[3] == Dict("b" => 1, "c" => 2, "d" => 3)
     @test Xenc.Column1 == [1.0, 2.0, 3.0, 4.0, 5.0]
-    @test Xenc.Column2 == [1.0, 2.0, 3.0, 4.0, 5.0]
-    @test Xenc.Column3 == [1, 2, 3]
+    @test Xenc.Column2 == Float32.([1.0, 2.0, 3.0, 4.0, 5.0])
+    @test Xenc.Column3 == Float32.([1, 2, 3])
     @test Xenc.Column4 == [1.0, 2.0, 3.0, 4.0, 5.0]
 
     X = coerce(X, :Column1 => Multiclass)
diff --git a/test/entity_embedding.jl b/test/entity_embedding.jl
index 738351a..da0c89b 100644
--- a/test/entity_embedding.jl
+++ b/test/entity_embedding.jl
@@ -1,13 +1,13 @@
 """
 See more functional tests in entity_embedding_utils.jl and mlj_model_interface.jl
 """
-
-batch = [
+batch = Float32.([
     0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 1.1;
-    1 2 3 4 5 6 7 8 9 10;
-    0.9 0.1 0.4 0.5 0.3 0.7 0.8 0.9 1.0 1.1
-    1 1 2 2 1 1 2 2 1 1
-]
+    1   2   3   4   5   6   7   8   9  10;
+    0.9 0.1 0.4 0.5 0.3 0.7 0.8 0.9 1.0 1.1;
+    1   1   2   2   1   1   2   2   1   1
+])
+
 
 entityprops = [
     (index = 2, levels = 10, newdim = 2),
@@ -145,7 +145,8 @@ end
     numfeats = 4
     embedder = MLJFlux.EntityEmbedder(entityprops, 4)
     output = embedder(batch)
-    @test output == batch
+    @test output ≈ batch
+    @test eltype(output) == Float32
 end
 
 
diff --git a/test/mlj_model_interface.jl b/test/mlj_model_interface.jl
index f7a0ee1..7984990 100644
--- a/test/mlj_model_interface.jl
+++ b/test/mlj_model_interface.jl
@@ -76,8 +76,8 @@ end
     nobservations = 12
     Xuser = rand(Float32, nobservations, 3)
     yuser = rand(Float32, nobservations)
-    alpha = rand(rng)
-    lambda = rand(rng)
+    alpha = rand(rng, Float32)
+    lambda = rand(rng, Float32)
     optimiser = Optimisers.Momentum()
     builder = MLJFlux.Linear()
     epochs = 1 # don't change this
@@ -94,7 +94,7 @@ end
         # (2) manually train for one epoch explicitly adding a loss penalty:
         chain = MLJFlux.build(builder, StableRNG(123), 3, 1);
         penalty = Penalizer(lambda, alpha); # defined in test_utils.jl
-        X, y = MLJFlux.collate(model, Xuser, yuser);
+        X, y = MLJFlux.collate(model, Xuser, yuser, 0);
         loss = model.loss;
         n_batches = div(nobservations, batch_size)
         optimiser_state = Optimisers.setup(optimiser, chain);
@@ -121,6 +121,7 @@ end
     # integration test:
     X, y = MLJBase.make_regression(10)
     X = Float32.(MLJBase.Tables.matrix(X)) |> MLJBase.Tables.table
+    y = Float32.(y)
     mach = MLJBase.machine(model, X, y)
     MLJBase.fit!(mach, verbosity=0)
     losses = MLJBase.training_losses(mach)
@@ -148,7 +149,8 @@ end
         builder = LisasBuilder(10),
     )
 
-    X, y = @load_boston
+    X = Tables.table(rand(Float32, 75, 2))
+    y = rand(Float32, 75)
     @test_logs(
         (:error, MLJFlux.ERR_BUILDER),
         @test_throws UndefVarError(:Chains) MLJBase.fit(model, 0, X, y)
@@ -165,16 +167,16 @@ end
     ]
     # table case
     X1 = (
-        Column1 = [1.0, 2.0, 3.0, 4.0, 5.0],
-        Column4 = [1.0, 2.0, 3.0, 4.0, 5.0],
-        Column5 = randn(5),
+        Column1 = Float32[1.0, 2.0, 3.0, 4.0, 5.0],
+        Column4 = Float32[1.0, 2.0, 3.0, 4.0, 5.0],
+        Column5 = randn(Float32, 5),
     )
     # matrix case
-    X2 = rand(5, 5)
+    X2 = rand(Float32, 5, 5)
     Xs = [X1, X2]
 
     y = categorical([0, 1, 0, 1, 1])
-    yreg = [0.1, -0.3, 0.2, 0.8, 0.9]
+    yreg = Float32[0.1, -0.3, 0.2, 0.8, 0.9]
     ys = [y, y, yreg, yreg]
     for j in eachindex(Xs)
         for i in eachindex(models)
@@ -210,15 +212,15 @@ end
     ]
 
     X = (
-        Column1 = [1.0, 2.0, 3.0, 4.0, 5.0],
+        Column1 = Float32[1.0, 2.0, 3.0, 4.0, 5.0],
         Column2 = categorical(['a', 'b', 'c', 'd', 'e']),
-        Column3 = [1.0, 2.0, 3.0, 4.0, 5.0],
-        Column4 = randn(5),
+        Column3 = Float32[1.0, 2.0, 3.0, 4.0, 5.0],
+        Column4 = randn(Float32, 5),
         Column5 = categorical(["group1", "group1", "group2", "group2", "group3"]),
     )
 
     y = categorical([0, 1, 0, 1, 1])
-    yreg = [0.1, -0.3, 0.2, 0.8, 0.9]
+    yreg = Float32[0.1, -0.3, 0.2, 0.8, 0.9]
     ys = [y, y, yreg, yreg]
 
     for i in eachindex(models)
@@ -272,15 +274,15 @@ end
     ]
 
     X = (
-        Column1 = [1.0, 2.0, 3.0, 4.0, 5.0],
+        Column1 = Float32[1.0, 2.0, 3.0, 4.0, 5.0],
         Column2 = categorical(['a', 'b', 'c', 'd', 'e']),
-        Column3 = [1.0, 2.0, 3.0, 4.0, 5.0],
-        Column4 = randn(5),
+        Column3 = Float32[1.0, 2.0, 3.0, 4.0, 5.0],
+        Column4 = randn(Float32, 5),
         Column5 = categorical(["group1", "group1", "group2", "group2", "group3"]),
     )
 
     y = categorical([0, 1, 0, 1, 1])
-    yreg = [0.1, -0.3, 0.2, 0.8, 0.9]
+    yreg = Float32[0.1, -0.3, 0.2, 0.8, 0.9]
     ys = [y, y, yreg, yreg]
 
     for i in eachindex(models)
@@ -316,7 +318,7 @@ end
             [:Column1, :Column2, :Column3, :Column4, :Column5],
         )
         @test mapping_matrices_fit != mapping_matrices_double_fit
-        # Try model prediction 
+        # Try model prediction
         Xpred = predict(mach, X)
     end
 end
diff --git a/test/regressor.jl b/test/regressor.jl
index 74c33a8..0b9c585 100644
--- a/test/regressor.jl
+++ b/test/regressor.jl
@@ -3,11 +3,11 @@ Random.seed!(123)
 N = 200
 Xm = MLJBase.table(randn(Float32, N, 5));   # purely numeric
 X = (; Tables.columntable(Xm)...,
-    Column1 = repeat([1.0, 2.0, 3.0, 4.0, 5.0], Int(N / 5)),
+    Column1 = repeat(Float32[1.0, 2.0, 3.0, 4.0, 5.0], Int(N / 5)),
     Column2 = categorical(repeat(['a', 'b', 'c', 'd', 'e'], Int(N / 5))),
     Column3 = categorical(repeat(["b", "c", "d", "f", "f"], Int(N / 5)), ordered = true),
-    Column4 = repeat([1.0, 2.0, 3.0, 4.0, 5.0], Int(N / 5)),
-    Column5 = randn(N),
+    Column4 = repeat(Float32[1.0, 2.0, 3.0, 4.0, 5.0], Int(N / 5)),
+    Column5 = randn(Float32, N),
     Column6 = categorical(
         repeat(["group1", "group1", "group2", "group2", "group3"], Int(N / 5)),
     ),
@@ -17,7 +17,7 @@ builder = MLJFlux.Short(σ = identity)
 optimiser = Optimisers.Adam()
 
 Random.seed!(123)
-y = 1 .+ X.x1 - X.x2 .- 2X.x4 + X.x5
+y = Float32(1) .+ X.x1 - X.x2 .- 2X.x4 + X.x5
 train, test = MLJBase.partition(1:N, 0.7)
 
 @testset_accelerated "NeuralNetworkRegressor" accel begin
diff --git a/test/runtests.jl b/test/runtests.jl
index e830f87..a6ed78e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -14,6 +14,7 @@ using StableRNGs
 using CUDA, cuDNN
 import StatisticalMeasures
 import Optimisers
+import Logging
 
 using ComputationalResources
 using ComputationalResources: CPU1, CUDALibs