diff --git a/src/encoders.jl b/src/encoders.jl
index de51714..60aac91 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(
+                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/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