fix serialization

JuliaAI · Jan 25, 2023 · 22b95d8 · 22b95d8
1 parent 58ebb8e
commit 22b95d8
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/src/MLJXGBoostInterface.jl b/src/MLJXGBoostInterface.jl
@@ -190,13 +190,9 @@ _save(fr; kw...) = XGB.save(fr, Vector{UInt8}; kw...)
 
 _restore(fr) = XGB.load(Booster, fr)
 
-MMI.save(::XGBoostAbstractRegressor, fr; kw...) = _save(fr; kw...)
+MMI.save(::XGTypes, fr; kw...) = (_save(fr[1]; kw...), fr[2])
 
-MMI.restore(::XGBoostAbstractRegressor, fr; kw...) = _restore(fr)
-
-MMI.save(::XGBoostClassifier, fr; kw...) = (_save(fr[1]; kw...), fr[2])
-
-MMI.restore(::XGBoostClassifier, fr) = (_restore(fr[1]), fr[2])
+MMI.restore(::XGTypes, fr) = (_restore(fr[1]), fr[2])
 
 MLJModelInterface.reports_feature_importances(::Type{<:XGBoostAbstractRegressor}) = true
 MLJModelInterface.reports_feature_importances(::Type{<:XGBoostAbstractClassifier}) = true