Merge pull request #428 from tidymodels/more-glmnet-cleanup

Handle single prediction case for multinomial glmnet

Author: topepo

NAMESPACE

@@ -283,4 +283,5 @@ importFrom(utils,getFromNamespace)
 importFrom(utils,globalVariables)
 importFrom(utils,head)
 importFrom(utils,methods)
+importFrom(vctrs,vec_size)
 importFrom(vctrs,vec_unique)

R/multinom_reg.R

@@ -168,13 +168,24 @@ check_args.multinom_reg <- function(object) {
 
 # ------------------------------------------------------------------------------
 
+#' @importFrom vctrs vec_size
 organize_multnet_class <- function(x, object) {
-  x[,1]
+  if (vec_size(x) > 1) {
+    x <- x[,1]
+  } else {
+    x <- as.character(x)
+  }
+  x
 }
 
+#' @importFrom vctrs vec_size
 organize_multnet_prob <- function(x, object) {
-  x <- x[,,1]
-  as_tibble(x)
+  if (vec_size(x) > 1) {
+    x <- as_tibble(x[,,1])
+  } else {
+    x <- tibble::as_tibble_row(x[,,1])
+  }
+  x
 }
 
 organize_nnet_prob <- function(x, object) {