Add tabnet_explain.model_fit

mlverse · Nov 3, 2022 · 2cab7d7 · 2cab7d7
1 parent cc78cc2
commit 2cab7d7
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/NAMESPACE b/NAMESPACE
@@ -4,6 +4,7 @@ S3method(predict,tabnet_fit)
 S3method(print,tabnet_fit)
 S3method(print,tabnet_pretrain)
 S3method(tabnet_explain,default)
+S3method(tabnet_explain,model_fit)
 S3method(tabnet_explain,tabnet_fit)
 S3method(tabnet_explain,tabnet_pretrain)
 S3method(tabnet_fit,data.frame)

diff --git a/R/explain.R b/R/explain.R
@@ -72,6 +72,11 @@ tabnet_explain.tabnet_fit <- function(object, new_data) {
 #' @rdname tabnet_explain
 tabnet_explain.tabnet_pretrain <- tabnet_explain.tabnet_fit
 
+#' @export
+#' @rdname tabnet_explain
+tabnet_explain.model_fit <- function(object, new_data) {
+  tabnet_explain(parsnip::extract_fit_engine(object), new_data)
+}
 
 convert_to_df <- function(x, nms) {
   x <- as.data.frame(as.matrix(x$to(device = "cpu")$detach()))

diff --git a/man/tabnet_explain.Rd b/man/tabnet_explain.Rd