diff --git a/Project.toml b/Project.toml
index e8e50b9..b477951 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "MLJIteration"
 uuid = "614be32b-d00c-4edb-bd02-1eb411ab5e55"
 authors = ["Anthony D. Blaom <anthony.blaom@gmail.com>"]
-version = "0.6.1"
+version = "0.6.2"
 
 [deps]
 IterationControl = "b3c1a2ee-3fec-4384-bf48-272ea71de57c"
@@ -11,7 +11,7 @@ Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 
 [compat]
 IterationControl = "0.5"
-MLJBase = "1.3"
+MLJBase = "1.4"
 julia = "1.6"
 
 [extras]
diff --git a/src/traits.jl b/src/traits.jl
index 28db5c6..6380d41 100644
--- a/src/traits.jl
+++ b/src/traits.jl
@@ -1,9 +1,7 @@
 MLJBase.is_wrapper(::Type{<:EitherIteratedModel}) = true
 MLJBase.caches_data_by_default(::Type{<:EitherIteratedModel}) = false
-MLJBase.load_path(::Type{<:DeterministicIteratedModel}) =
-    "MLJIteration.DeterministicIteratedModel"
-MLJBase.load_path(::Type{<:ProbabilisticIteratedModel}) =
-    "MLJIteration.ProbabilisticIteratedModel"
+MLJBase.load_path(::Type{<:EitherIteratedModel}) = "MLJIteration.IteratedModel"
+MLJBase.constructor(::Type{<:EitherIteratedModel}) = IteratedModel
 MLJBase.package_name(::Type{<:EitherIteratedModel}) = "MLJIteration"
 MLJBase.package_uuid(::Type{<:EitherIteratedModel}) =
     "614be32b-d00c-4edb-bd02-1eb411ab5e55"
diff --git a/test/traits.jl b/test/traits.jl
index b63f03f..341628b 100644
--- a/test/traits.jl
+++ b/test/traits.jl
@@ -13,7 +13,7 @@ imodel = IteratedModel(model=model, measure=mae)
 @test !MLJBase.caches_data_by_default(imodel)
 @test !supports_weights(imodel)
 @test !supports_class_weights(imodel)
-@test load_path(imodel) == "MLJIteration.DeterministicIteratedModel"
+@test load_path(imodel) == "MLJIteration.IteratedModel"
 @test package_name(imodel) == "MLJIteration"
 @test package_uuid(imodel) == "614be32b-d00c-4edb-bd02-1eb411ab5e55"
 @test package_url(imodel) == "https://github.com/JuliaAI/MLJIteration.jl"
@@ -22,6 +22,7 @@ imodel = IteratedModel(model=model, measure=mae)
 @test input_scitype(imodel) == input_scitype(model)
 @test output_scitype(imodel) == output_scitype(model)
 @test target_scitype(imodel) == target_scitype(model)
+@test constructor(imodel) == IteratedModel
 
 end