diff --git a/.Rbuildignore b/.Rbuildignore
index 4117da74..15030367 100644
--- a/.Rbuildignore
+++ b/.Rbuildignore
@@ -14,3 +14,5 @@
 ^CRAN-SUBMISSION$
 ^revdep$
 ^vignettes/*_files$
+^\.claude$
+^\.positai$
diff --git a/.gitignore b/.gitignore
index 4b3d4c83..16879333 100644
--- a/.gitignore
+++ b/.gitignore
@@ -15,3 +15,5 @@ tabnet_*.tar.gz
 tabnet.Rproj
 po/glossary.csv
 inst/IMPORTLIST
+.positai
+tools
diff --git a/DESCRIPTION b/DESCRIPTION
index 95044dcc..d8976843 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -46,6 +46,7 @@ Imports:
     zeallot
 Suggests: 
     cli,
+    fBasics,
     knitr,
     modeldata,
     patchwork,
@@ -68,5 +69,5 @@ Config/testthat/parallel: false
 Config/testthat/start-first: interface, explain, params
 Encoding: UTF-8
 Roxygen: list(markdown = TRUE)
-RoxygenNote: 7.3.3
 Language: en-US
+Config/roxygen2/version: 8.0.0
diff --git a/NAMESPACE b/NAMESPACE
index 8f99e229..0bb1617c 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -24,6 +24,7 @@ S3method(tabnet_pretrain,recipe)
 S3method(update,tabnet)
 export("%>%")
 export(attention_width)
+export(build_ancestor_matrix_from_outcomes)
 export(cat_emb_dim)
 export(check_compliant_node)
 export(checkpoint_epochs)
@@ -39,6 +40,9 @@ export(mlp_activation)
 export(mlp_hidden_multiplier)
 export(momentum)
 export(nn_aum_loss)
+export(nn_mc_loss)
+export(nnf_mc_loss)
+export(nnf_multilabel_one_hot)
 export(node_to_df)
 export(num_independent)
 export(num_independent_decoder)
@@ -71,6 +75,18 @@ importFrom(rlang,.data)
 importFrom(stats,predict)
 importFrom(stats,update)
 importFrom(tidyr,replace_na)
+importFrom(torch,as_array)
+importFrom(torch,nn_module)
 importFrom(torch,nn_prune_head)
+importFrom(torch,torch_argsort)
+importFrom(torch,torch_cat)
+importFrom(torch,torch_long)
+importFrom(torch,torch_matmul)
+importFrom(torch,torch_mean)
+importFrom(torch,torch_minimum)
+importFrom(torch,torch_mul)
+importFrom(torch,torch_std)
+importFrom(torch,torch_sum)
+importFrom(torch,torch_tensor)
 importFrom(tune,min_grid)
 importFrom(zeallot,"%<-%")
diff --git a/NEWS.md b/NEWS.md
index 88951d32..1c5b1f25 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,5 +1,9 @@
 # tabnet (development version)
 
+## Bugfixes
+
+* Ancestor matrix is now taken into account for hierarchical classification (#188).
+
 # tabnet 0.8.0
 
 ## New features
diff --git a/R/hardhat.R b/R/hardhat.R
index 1d57a1b4..456d0a14 100644
--- a/R/hardhat.R
+++ b/R/hardhat.R
@@ -162,16 +162,14 @@ tabnet_fit.Node <- function(x, tabnet_model = NULL, config = tabnet_config(), ..
   # get tree leaves and extract attributes into data.frames
   xy_df <- node_to_df(x)
   processed <- hardhat::mold(xy_df$x, xy_df$y)
-  # Given n classes, M is an (n x n) matrix where M_ij = 1 if class i is descendant of class j
-  ancestor <- data.tree::ToDataFrameNetwork(x) %>%
-   mutate_if(is.character, ~.x %>% as.factor %>% as.numeric)
-  # TODO check correctness
-  # embed the M matrix in the config$ancestor variable
-  dims <- c(max(ancestor), max(ancestor))
-  ancestor_m <- Matrix::sparseMatrix(ancestor$from, ancestor$to, dims = dims, x = 1)
   check_type(processed$outcomes)
-
+  
   config <- merge_config_and_dots(config, ...)
+  # add ancestor boolean sparse matrix to config
+  # check_dag_compliance(xy_df$y)
+  config$ancestor <- build_ancestor_matrix_from_outcomes(x, processed$outcomes)
+  # make outcomes levels available so that batched y could be one-hot encoded.
+  config$outcomes <- processed$outcomes
   tabnet_bridge(processed, config = config, tabnet_model, from_epoch, task = "supervised")
 }
 
@@ -272,7 +270,7 @@ tabnet_pretrain.default <- function(x, ...) {
 
 #' @export
 #' @rdname tabnet_pretrain
-tabnet_pretrain.data.frame <- function(x, y, tabnet_model = NULL, config = tabnet_config(), ..., from_epoch = NULL) {
+tabnet_pretrain.data.frame <- function(x, y = NULL, tabnet_model = NULL, config = tabnet_config(), ..., from_epoch = NULL) {
   processed <- hardhat::mold(x, y)
 
   config <- merge_config_and_dots(config, ...)
@@ -309,8 +307,7 @@ tabnet_pretrain.Node <- function(x, tabnet_model = NULL, config = tabnet_config(
   check_compliant_node(x)
   # get tree leaves and extract attributes into data.frames
   xy_df <- node_to_df(x)
-  tabnet_pretrain(xy_df$x, xy_df$y, tabnet_model = tabnet_model, config = config, ..., from_epoch = from_epoch)
-
+  tabnet_pretrain(xy_df$x, tabnet_model = tabnet_model, config = config, ..., from_epoch = from_epoch)
 }
 
 new_tabnet_pretrain <- function(pretrain, blueprint) {
@@ -418,13 +415,15 @@ tabnet_bridge <- function(processed, config = tabnet_config(), tabnet_model, fro
 #' @importFrom stats predict
 #' @export
 predict.tabnet_fit <- function(object, new_data, type = NULL, ..., epoch = NULL) {
-  if (inherits(new_data, "Node")) {
+  if (inherits(new_data, "Node") && !is.null(object$fit$config$ancestor)) {
     new_data_df <- node_to_df(new_data)$x
+    # Enforces column order, type, column names, etc
+    processed <- hardhat::forge(new_data_df, object$blueprint)
+    
   } else {
     new_data_df <- new_data
+    processed <- hardhat::forge(new_data, object$blueprint)
   }
-  # Enforces column order, type, column names, etc
-  processed <- hardhat::forge(new_data_df, object$blueprint)
   batch_size <- object$fit$config$batch_size
   out <- predict_tabnet_bridge(type, object, processed$predictors, epoch, batch_size)
   hardhat::validate_prediction_size(out, new_data_df)
@@ -436,8 +435,8 @@ predict_tabnet_bridge <- function(type, object, predictors, epoch, batch_size) {
   type <- check_type(object$blueprint$ptypes$outcomes, type)
   is_multi_outcome <- ncol(object$blueprint$ptypes$outcomes) > 1
   outcome_nlevels <- NULL
-  if (is_multi_outcome & type != "numeric") {
-    outcome_nlevels <- purrr::map_dbl(object$blueprint$ptypes$outcomes, ~length(levels(.x)))
+  if (is_multi_outcome && type != "numeric") {
+    outcome_nlevels <- purrr::map_dbl(object$blueprint$ptypes$outcomes, ~nlevels(.x))
   }
 
   if (!is.null(epoch)) {
@@ -458,6 +457,7 @@ predict_tabnet_bridge <- function(type, object, predictors, epoch, batch_size) {
     object$fit$network$load_state_dict(m$state_dict())
   }
 
+  
   type_multioutcome <- paste0(type, "_", is_multi_outcome)
   switch(
     type_multioutcome,
@@ -604,4 +604,81 @@ nn_prune_head.tabnet_pretrain <- function(x, head_size) {
     nn_prune_head(x$fit$network, head_size=head_size)
   }
 
-}
\ No newline at end of file
+}
+
+
+#' Build ancestor matrix aligned with observed outcome classes
+#'
+#' Extracts class names from the outcome tibble (factor levels) and builds
+#' the ancestor matrix only for classes that actually appear in the data.
+#'
+#' @param x A `data.tree::Node` object.
+#' @param outcomes A tibble with factor columns (one per hierarchy level),
+#'   as returned by `hardhat::mold()$outcomes`.
+#' @param device Torch device ("cpu" or "cuda").
+#' @return A `torch_tensor` of shape `(1, n_classes, n_classes)`.
+#' @export
+build_ancestor_matrix_from_outcomes <- function(x, outcomes, device = "cpu") {
+  # 1. Extract all class names from factor levels (preserving order)
+  #    outcomes is a tibble with one factor column per hierarchy level
+  level_cols <- names(outcomes)
+  all_class_names <- unlist(lapply(outcomes, levels), use.names = FALSE)
+  n_classes <- length(all_class_names)
+  
+  if (n_classes == 0L) {
+    runtime_error("No factor levels found in outcomes : {str(outcomes)}")
+  }
+  
+  # 2. Build a lookup: class_name -> data.tree Node
+  all_nodes <- data.tree::Traverse(x, traversal = "pre-order")
+  all_nodes <- unname(all_nodes)
+  level_lengths <- lengths(lapply(outcomes, levels))
+  lvl_vector <- rep(seq_along(level_cols) + 1L, level_lengths)
+  
+  # 3. Resolve each class name to its Node
+  class_nodes <- lapply(seq_along(all_class_names), function(k) {
+    nm <- all_class_names[k]
+    lvl <- lvl_vector[k]
+    
+    candidates <- Filter(function(n) n$level == lvl && n$name == nm, all_nodes)
+    if (length(candidates) == 0) {
+      runtime_error("Factor level {.var {nm}} not found at tree level {lvl} (outcomes column {.var {level_cols[lvl - 1L]}})")
+    }
+    candidates[[1]]
+  })
+  
+  # 4. Create 1-based index mapping
+  class_map <- setNames(seq_len(n_classes), all_class_names)
+  
+  # 5. Collect (descendant, ancestor) pairs by climbing up
+  row_list <- vector("list", n_classes)
+  col_list <- vector("list", n_classes)
+  
+  for (i in seq_len(n_classes)) {
+    current <- class_nodes[[i]]
+    anc_indices <- integer()
+    
+    repeat {
+      idx <- class_map[current$name]
+      if (!is.null(idx)) {
+        anc_indices <- c(anc_indices, idx)
+      }
+      if (current$isRoot || is.null(current$parent)) break
+      current <- current$parent
+    }
+    
+    row_list[[i]] <- rep(i, length(anc_indices))
+    col_list[[i]] <- anc_indices
+  }
+  
+  # 6. Fill matrix
+  R <- matrix(0L, nrow = n_classes, ncol = n_classes)
+  rows <- unlist(row_list, use.names = FALSE)
+  cols <- unlist(col_list, use.names = FALSE)
+  if (length(rows) > 0) R[cbind(rows, cols)] <- 1L
+  
+  # 7. Convert to torch
+  R_torch <- torch::torch_tensor(R, dtype = torch::torch_double(), device = device)
+  R_torch$unsqueeze(1)
+}
+
diff --git a/R/loss.R b/R/loss.R
index ef018dc3..89eed651 100644
--- a/R/loss.R
+++ b/R/loss.R
@@ -4,7 +4,8 @@
 #' element in the input \eqn{y_pred} and target \eqn{embedded_x} on the values masked by \eqn{obfuscation_mask}. 
 #' 
 #' @noRd
-nn_unsupervised_loss <- torch::nn_module(
+#' @importFrom torch torch_mul torch_std torch_matmul torch_mean
+nn_unsupervised_loss <- nn_module(
   "nn_unsupervised_loss",
   inherit = torch::nn_cross_entropy_loss,
   
@@ -15,15 +16,15 @@ nn_unsupervised_loss <- torch::nn_module(
   
   forward = function(y_pred, embedded_x, obfuscation_mask){
     errors <- y_pred - embedded_x
-    reconstruction_errors <- torch::torch_mul(errors, obfuscation_mask) ^ 2
-    batch_stds <- torch::torch_std(embedded_x, dim = 1) ^ 2 + self$eps
+    reconstruction_errors <- torch_mul(errors, obfuscation_mask) ^ 2
+    batch_stds <- torch_std(embedded_x, dim = 1) ^ 2 + self$eps
     
     # compute the number of obfuscated variables to reconstruct
-    nb_reconstructed_variables <- torch::torch_sum(obfuscation_mask, dim = 2)
+    nb_reconstructed_variables <- torch_sum(obfuscation_mask, dim = 2)
     
     # take the mean of the reconstructed variable errors
-    features_loss <- torch::torch_matmul(reconstruction_errors, 1 / batch_stds) / (nb_reconstructed_variables +  self$eps)
-    loss <- torch::torch_mean(features_loss, dim = 1)
+    features_loss <- torch_matmul(reconstruction_errors, 1 / batch_stds) / (nb_reconstructed_variables +  self$eps)
+    loss <- torch_mean(features_loss, dim = 1)
     loss
   }
 )
@@ -45,7 +46,9 @@ nn_unsupervised_loss <- torch::nn_module(
 #' output <- loss(input, target)
 #' output$backward()
 #' @export
-nn_aum_loss <- torch::nn_module(
+#' @importFrom torch nn_module torch_sum torch_cat torch_minimum torch_long torch_argsort
+#' @importFrom torch torch_tensor as_array
+nn_aum_loss <- nn_module(
   "nn_aum_loss",
   inherit = torch::nn_mse_loss,
   initialize = function(){
@@ -57,8 +60,8 @@ nn_aum_loss <- torch::nn_module(
     is_positive <- label_tensor == label_tensor$max()
     is_negative <- is_positive$bitwise_not()
     # manage case when prediction error is null (prevent division by 0)
-    if(as.logical(torch::torch_sum(is_positive) == 0) || as.logical(torch::torch_sum(is_negative) == 0)){
-      return(torch::torch_sum(pred_tensor*0))
+    if(as.logical(torch_sum(is_positive) == 0) || as.logical(torch_sum(is_negative) == 0)){
+      return(torch_sum(pred_tensor*0))
     }
     
     # pred tensor may be [prediction, case_wts] when add_case_weight() is used. We keep only prediction
@@ -68,10 +71,10 @@ nn_aum_loss <- torch::nn_module(
     
     # nominal case
     fn_diff <- -1L * is_positive
-    fp_diff <- is_negative$to(dtype = torch::torch_long())
-    fp_denom <- torch::torch_sum(is_negative) # or 1 for AUM based on count instead of rate
-    fn_denom <- torch::torch_sum(is_positive) # or 1 for AUM based on count instead of rate
-    sorted_pred_ids <- torch::torch_argsort(pred_tensor, dim = 1, descending = TRUE)$squeeze(-1)
+    fp_diff <- is_negative$to(dtype = torch_long())
+    fp_denom <- torch_sum(is_negative) # or 1 for AUM based on count instead of rate
+    fn_denom <- torch_sum(is_positive) # or 1 for AUM based on count instead of rate
+    sorted_pred_ids <- torch_argsort(pred_tensor, dim = 1, descending = TRUE)$squeeze(-1)
     
     sorted_fp_cum <- fp_diff[sorted_pred_ids]$cumsum(dim = 1) / fp_denom
     sorted_fn_cum <- -fn_diff[sorted_pred_ids]$flip(1)$cumsum(dim = 1)$flip(1) / fn_denom
@@ -79,26 +82,259 @@ nn_aum_loss <- torch::nn_module(
     sorted_dedup <- sorted_thresh_gr$diff(dim = 1) != 0
     # pad to replace removed last element
     padding <- sorted_dedup$slice(dim = 1, 0, 1) # torch_tensor 1 w same dtype, same shape, same device 
-    sorted_fp_end <- torch::torch_cat(c(sorted_dedup, padding))
-    sorted_fn_end <- torch::torch_cat(c(padding, sorted_dedup))
+    sorted_fp_end <- torch_cat(c(sorted_dedup, padding))
+    sorted_fn_end <- torch_cat(c(padding, sorted_dedup))
     uniq_thresh_gr <- sorted_thresh_gr[sorted_fp_end]
     uniq_fp_after <- sorted_fp_cum[sorted_fp_end]
     uniq_fn_before <- sorted_fn_cum[sorted_fn_end]
     if (pred_tensor$ndim == 1) {
-      FPR <- torch::torch_cat(c(padding$logical_not(), uniq_fp_after)) # FPR with trailing 0
-      FNR <-  torch::torch_cat(c(uniq_fn_before, padding$logical_not())) # FNR with leading 0
+      FPR <- torch_cat(c(padding$logical_not(), uniq_fp_after)) # FPR with trailing 0
+      FNR <-  torch_cat(c(uniq_fn_before, padding$logical_not())) # FNR with leading 0
       self$roc_aum <- list(
         FPR = FPR,
         FNR = FNR,
         TPR = 1 - FNR,
-        "min(FPR,FNR)" = torch::torch_minimum(FNR, FPR), # full-range min(FNR, FPR)
-        constant_range_low = torch::torch_cat(c(torch::torch_tensor(-Inf), uniq_thresh_gr)),
-        constant_range_high = torch::torch_cat(c(uniq_thresh_gr, torch::torch_tensor(Inf)))
-      ) %>% purrr::map_dfc(torch::as_array)
+        "min(FPR,FNR)" = torch_minimum(FNR, FPR), # full-range min(FNR, FPR)
+        constant_range_low = torch_cat(c(torch_tensor(-Inf), uniq_thresh_gr)),
+        constant_range_high = torch_cat(c(uniq_thresh_gr, torch_tensor(Inf)))
+      ) %>% purrr::map_dfc(as_array)
     }
-    min_FPR_FNR <- torch::torch_minimum(uniq_fp_after[1:-2], uniq_fn_before[2:N])
+    min_FPR_FNR <- torch_minimum(uniq_fp_after[1:-2], uniq_fn_before[2:N])
     constant_range_gr <- uniq_thresh_gr$diff() # range splits leading to {FPR, FNR } errors (see roc_aum row)
-    torch::torch_sum(min_FPR_FNR * constant_range_gr, dim = 1)
+    torch_sum(min_FPR_FNR * constant_range_gr, dim = 1)
     
   }
 )
+
+
+#' Apply hierarchy constraints via max-pooling over descendants (MCM)
+#'
+#' Given neural network outputs x and ancestor matrix R, enforces that
+#' if a class is predicted positive, all its ancestors must also be positive.
+#' Implements: `final_out[i] = max{x[j] : R[i,j] = 1}`
+#'
+#' @param x A `torch_tensor` of shape `(batch_size, n_classes)`.
+#' @param R A `torch_tensor` of shape `(1, n_classes, n_classes)` where 
+#'   `R[1, i, j] = 1` iff class `i` is a descendant of class `j`.
+#' @return A `torch_tensor` of shape `(batch_size, n_classes)` with constrained outputs.
+get_constr_output <- function(x, R) {
+  c_out <- x$double()$unsqueeze(2)$expand(c(x$shape[1], R$shape[2], R$shape[2]))
+  R_batch <- R$expand(c(x$shape[1], R$shape[2], R$shape[2]))
+  final_out <- (R_batch * c_out)$clone()$max(dim = 3)
+  final_out[[1]]
+}
+
+
+#' Max-Constraint Margin Loss (functional)
+#'
+#' Computes the hierarchy-constrained loss for multi-label classification.
+#' Enforces that if a class is predicted positive, all its ancestors must 
+#' also be positive, using the ancestor matrix R.
+#'
+#' The loss combines constrained outputs differently for positive and negative 
+#' labels:
+#' \itemize{
+#'   \item For positive labels: uses constrained output of label-weighted predictions
+#'   \item For negative labels: uses constrained raw predictions (penalizes ancestor violations)
+#' }
+#'
+#' @param output A `torch_tensor` of raw network outputs (pre-sigmoid), 
+#'   shape `(batch_size, n_classes)`.
+#' @param target Binary target labels, shape `(batch_size, n_classes)`.
+#' @param R Ancestor matrix tensor of shape `(1, n_classes, n_classes)` where 
+#'   `R[1, i, j] = 1` iff class `i` is a descendant of class `j`.
+#' @param to_eval Optional logical tensor of shape `(n_classes,)` indicating 
+#'   which classes to include in the loss computation. If `NULL`, all classes 
+#'   are evaluated.
+#' @param criterion Loss function to apply after constraint propagation. 
+#'   Default: `nnf_binary_cross_entropy_with_logits` (expects raw logits).
+#'
+#' @return A scalar `torch_tensor` containing the computed loss, or a tensor 
+#'   of shape `(batch_size, n_classes)` if `reduction = "none"`.
+#'
+#' @seealso [nn_mc_loss()], [get_constr_output()]
+#' @export
+nnf_mc_loss <- function(output, target, R, to_eval = NULL, 
+                        criterion = nnf_binary_cross_entropy_with_logits) {
+  # Ensure double precision for numerical stability during constraint propagation
+  output_d <- output$double()
+  
+  # 1. Constrained output from raw predictions: max-pool over descendants
+  constr_output <- get_constr_output(output_d, R)  # (batch, n_classes)
+  
+  # 2. Label-weighted output, then constrained (for positive label handling)
+  labeled_output <- target * output_d
+  train_output <- get_constr_output(labeled_output, R)
+  
+  # 3. Blend outputs based on ground-truth labels:
+  #    - Positive labels: use constrained label-weighted output
+  #    - Negative labels: use constrained raw output
+  blended_output <- (1 - target) * constr_output + target * train_output
+  
+  # 4. Select classes to evaluate (if specified)
+  if (!is.null(to_eval)) {
+    blended_output <- blended_output[, to_eval, drop = FALSE]
+    target <- target[, to_eval, drop = FALSE]
+  }
+  
+  # 5. Apply the base loss function (e.g., BCE with logits)
+  loss <- criterion(
+    blended_output, 
+    target$double()
+  )
+  
+  return(loss)
+}
+
+
+#' Max-Constraint Margin Loss (module)
+#'
+#' Module wrapper for [nnf_mc_loss()] with configurable parameters.
+#' Stores the ancestor matrix R and evaluation mask for reuse across batches.
+#'
+#' @param R Ancestor matrix tensor of shape `(1, n_classes, n_classes)`.
+#' @param to_eval Optional logical tensor of shape `(n_classes,)` indicating 
+#'   which classes to include in loss computation.
+#' @param criterion Loss function module or functional to apply after constraint 
+#'   propagation. Default: `nn_binary_cross_entropy_with_logits()`.
+#' @param reduction (string, optional): Reduction method: `'none'` | `'mean'` | `'sum'`.
+#'
+#' @section Shape:
+#' - Input `output`: \eqn{(N, C)} where N = batch size, C = number of classes
+#' - Input `target`: \eqn{(N, C)}, same shape as output, binary values
+#' - Output: scalar by default. If `reduction = "none"`, then \eqn{(N, C')} 
+#'   where C' is the number of evaluated classes
+#'
+#' @examples
+#' \dontrun{
+#' # Build ancestor matrix from hierarchy
+#' R <- build_ancestor_matrix_from_outcomes(my_tree, processed$outcomes, device = "cuda")
+#' 
+#' # Create loss module
+#' loss_fn <- nn_mc_loss(R = R, reduction = "mean")
+#' 
+#' # Forward pass
+#' output <- model(x)  # (batch, n_classes)
+#' loss <- loss_fn(output, labels)
+#' loss$backward()
+#' }
+#'
+#' @seealso [nnf_mc_loss()], [build_ancestor_matrix_from_outcomes()], [get_constr_output()]
+#' @export
+nn_mc_loss <- nn_module(
+  "nn_mc_loss",
+  inherit = torch::nn_l1_loss,
+  
+  initialize = function(R, to_eval = NULL, 
+                        criterion = torch::nnf_binary_cross_entropy_with_logits,
+                        reduction = "mean") {
+    super$initialize(reduction = reduction)
+    
+    # Store ancestor matrix (move to device if needed)
+    self$R <- R
+    self$to_eval <- to_eval
+    # Resolve criterion based on its type
+    self$criterion_fn <- .resolve_mc_criterion(criterion, reduction)
+  },
+  
+  forward = function(output, target) {
+    nnf_mc_loss(
+      output = output,
+      target = target,
+      R = self$R,
+      to_eval = self$to_eval,
+      criterion = self$criterion_fn
+    )
+  }
+)
+
+#' Resolve criterion into a callable function(input, target, reduction)
+#' @keywords internal
+#' @noRd
+.resolve_mc_criterion <- function(criterion, reduction) {
+  # Case 1: Already an nn_module instance
+  if (inherits(criterion, "nn_module")) {
+    module_reduction <- criterion$reduction
+    if (!is.null(module_reduction) && module_reduction != reduction) {
+      warn(
+        c(
+          "The criterion module has reduction={.val {module_reduction}}",
+          "but nn_mc_loss was called with reduction={.val {reduction}}.",
+          "i" = "The module's reduction will be used."
+        ),
+        class = "mc_loss_reduction_mismatch"
+      )
+    }
+    return(function(input, target) criterion(input, target))
+  }
+  
+  # Case 2: A function (could be functional nnf_* or constructor nn_*)
+  if (rlang::is_function(criterion)) {
+    # Try to detect if it's a constructor by calling with just reduction
+    # Constructors return nn_module, functionals need input/target
+    maybe_module <- tryCatch(
+      {
+        result <- criterion(reduction = reduction)
+        if (inherits(result, "nn_module")) result else NULL
+      },
+      error = function(e) NULL
+    )
+    
+    if (!is.null(maybe_module)) {
+      # It's a constructor (e.g., nn_bce_with_logits_loss)
+      return(function(input, target) maybe_module(input, target))
+    }
+    
+    # It's a functional (e.g., nnf_binary_cross_entropy_with_logits)
+    return(function(input, target) {
+      criterion(input, target, reduction = reduction)
+    })
+  }
+  
+  # Invalid type
+  value_error(
+    c(
+      "`criterion` must be a function or an `nn_module`.",
+      "x" = "Got: {.class {class(criterion)[1]}}"
+    ),
+    class = "mc_loss_invalid_criterion"
+  )
+}
+
+#' Convert class_id tensor to binary one-hot tensor
+#'
+#' Transforms a tensor of class indices (one column per hierarchy level)
+#' into a binary tensor where each column corresponds to a class.
+#'
+#' @param y A `torch_tensor` of shape `(batch_size, n_levels)` containing
+#'   1-based class indices.
+#' @param outcomes A tibble with factor columns (as from `hardhat::mold()$outcomes`).
+#' @param device Torch device.
+#' @return A `torch_tensor` of shape `(batch_size, n_classes)` with binary values.
+#' @export
+nnf_multilabel_one_hot <- function(y, outcomes, device = "cpu") {
+  batch_size <- y$shape[1]
+  n_levels <- y$shape[2]
+  
+  # Number of classes per level
+  n_per_level <- lengths(lapply(outcomes, levels))
+  n_classes <- sum(n_per_level)
+  
+  one_hot_list <- vector("list", n_levels)
+  
+  for (lvl in seq_len(n_levels)) {
+    level_ids <- y[, lvl]$to(dtype = torch::torch_long())
+    
+    # Encode one-hot of each levels
+    one_hot_list[[lvl]] <- torch::nnf_one_hot(
+      level_ids, 
+      num_classes = n_per_level[lvl]
+    )
+  }
+  # concatenate along the columns axis)
+  torch::torch_cat(one_hot_list, dim = 2)$to(
+    dtype = torch::torch_double(), 
+    device = device
+  )
+}
+
diff --git a/R/explain.R b/R/model_explain.R
similarity index 100%
rename from R/explain.R
rename to R/model_explain.R
diff --git a/R/pretraining.R b/R/model_pretraining.R
similarity index 99%
rename from R/pretraining.R
rename to R/model_pretraining.R
index 2ec315d6..0f48482e 100644
--- a/R/pretraining.R
+++ b/R/model_pretraining.R
@@ -178,9 +178,9 @@ tabnet_train_unsupervised <- function(x, config = tabnet_config(), epoch_shift =
       metrics[[epoch]][["valid"]] <- transpose_metrics(valid_metrics)$loss
     }
 
-    if (config$verbose & !has_valid)
+    if (config$verbose && !has_valid)
       message(gettextf("[Epoch %03d] Loss: %3f", epoch, mean(metrics[[epoch]]$train)))
-    if (config$verbose & has_valid)
+    if (config$verbose && has_valid)
       message(gettextf("[Epoch %03d] Loss: %3f, Valid loss: %3f", epoch, mean(metrics[[epoch]]$train), mean(metrics[[epoch]]$valid)))
 
     # Early-stopping checks
diff --git a/R/model.R b/R/model_training.R
similarity index 89%
rename from R/model.R
rename to R/model_training.R
index 30084761..f57984a6 100644
--- a/R/model.R
+++ b/R/model_training.R
@@ -175,7 +175,8 @@ tabnet_config <- function(batch_size = 1024^2,
                           early_stopping_tolerance = 0,
                           early_stopping_patience = 0L,
                           num_workers=0L,
-                          skip_importance = FALSE) {
+                          skip_importance = FALSE
+                          ) {
   if (is.null(decision_width) && is.null(attention_width)) {
     decision_width <- 8 # default is 8
   }
@@ -226,20 +227,6 @@ tabnet_config <- function(batch_size = 1024^2,
   )
 }
 
-get_constr_output <- function(x, R) {
-    # MCM of the prediction given the hierarchy constraint expressed in the matrix R """
-    c_out <- x$unsqueeze(2)$expand(c(x$shape[1], R$shape[2], R$shape[2]))
-    R_batch <- R$expand(c(x$shape[1], R$shape[2], R$shape[2]))
-    final_out <- torch::torch_max(R_batch * c_out, dim = 3)
-    final_out[[1]]
-}
-
-max_constraint_output <- function(output, labels, ancestor) {
-  constr_output <-  get_constr_output(output, ancestor)
-  train_output <-  get_constr_output(labels * output, ancestor)
-  labels$bitwise_not() * constr_output + labels * train_output
-}
-
 resolve_loss <- function(config, dtype) {
   loss <- config$loss
 
@@ -249,7 +236,9 @@ resolve_loss <- function(config, dtype) {
     loss_fn <- loss
   else if (loss %in% c("mse", "auto") && !dtype == torch::torch_long())
     loss_fn <- torch::nn_mse_loss()
-  else if ((loss %in% c("bce", "cross_entropy", "auto") && dtype == torch::torch_long()) || !is.null(config$ancestor_tt))
+  else if (!is.null(config$ancestor))
+    loss_fn <- nn_mc_loss(R = config$ancestor)
+  else if ((loss %in% c("bce", "cross_entropy", "auto") && dtype == torch::torch_long()))
     # cross entropy loss is required
     loss_fn <- torch::nn_cross_entropy_loss()
   else
@@ -270,42 +259,37 @@ resolve_early_stop_monitor <- function(early_stopping_monitor, valid_split) {
 }
 
 train_batch <- function(network, optimizer, batch, config) {
-  # NULLing values to avoid a R-CMD Check Note "No visible binding for global variable"
+  # NULL-ing values to avoid a R-CMD Check Note "No visible binding for global variable"
   out <- M_loss <- NULL
   # forward pass
   c(out, M_loss) %<-% network(batch$x, batch$x_na_mask)
-  # if target is multi-outcome, loss has to be applied to each label-group
-  if (max(batch$output_dim$shape) > 1) {
-    # multi-outcome
+
+  # if target is multi-outcome but not max_constraint loss, loss has to be applied to each label-group
+  if (max(batch$output_dim$shape) > 1 && is.null(config$ancestor)) {
+    # standard multi-outcome
     outcome_nlevels <- as.numeric(batch$output_dim$to(device="cpu"))
-    if (!is.null(config$ancestor_tt)) {
-      # hierarchical mandates use of `max_constraint_output`
-      loss <- torch::torch_sum(torch::torch_stack(purrr::pmap(
-        list(
-          torch::torch_split(out, outcome_nlevels, dim = 2),
-          torch::torch_split(batch$y, rep(1, length(outcome_nlevels)), dim = 2)
-        ),
-        ~config$loss_fn(max_constraint_output(.x, .y$squeeze(2), config$ancestor_tt))
-      )),
-      dim = 1)
-    } else {
-      # use `resolved_loss`
-      loss <- torch::torch_sum(torch::torch_stack(purrr::pmap(
-        list(
-          torch::torch_split(out, outcome_nlevels, dim = 2),
-          torch::torch_split(batch$y, rep(1, length(outcome_nlevels)), dim = 2)
-        ),
-        ~config$loss_fn(.x, .y$squeeze(2))
-      )),
-      dim = 1)
-    }
+
+    # use `resolved_loss`
+    loss <- torch::torch_sum(torch::torch_stack(purrr::pmap(
+      list(
+        torch::torch_split(out, outcome_nlevels, dim = 2),
+        torch::torch_split(batch$y, rep(1, length(outcome_nlevels)), dim = 2)
+      ),
+      ~config$loss_fn(.x, .y$squeeze(2))
+    )),
+    dim = 1)
+  } else if (!is.null(config$ancestor)) {
+    # multi-outcome max_constraint loss ned one-hot encoding of targets
+    loss <- config$loss_fn(out, nnf_multilabel_one_hot(
+      y = batch$y,
+      outcomes = config$outcomes,
+      device = out$device
+    ))
+  } else if (batch$y$dtype == torch::torch_long()) {
+    # classifier needs a squeeze for bce loss
+    loss <- config$loss_fn(out, batch$y$squeeze(2))
   } else {
-    if (batch$y$dtype == torch::torch_long()) {
-      # classifier needs a squeeze for bce loss
-      loss <- config$loss_fn(out, batch$y$squeeze(2))
-    } else {
-      loss <- config$loss_fn(out, batch$y)
-    }
+    loss <- config$loss_fn(out, batch$y)
   }
   # Add the overall sparsity loss
   loss <- loss - config$lambda_sparse * M_loss
@@ -329,30 +313,26 @@ valid_batch <- function(network, batch, config) {
   # forward pass
   c(out, M_loss) %<-% network(batch$x, batch$x_na_mask)
   # loss has to be applied to each label-group when output_dim is a vector
-  if (max(batch$output_dim$shape) > 1) {
-    # multi-outcome
+  if (max(batch$output_dim$shape) > 1 && is.null(config$ancestor)) {
+    # standard multi-outcome
     outcome_nlevels <- as.numeric(batch$output_dim$to(device="cpu"))
-    if (!is.null(config$ancestor_tt)) {
-      # hierarchical mandates use of `max_constraint_output`
-      loss <- torch::torch_sum(torch::torch_stack(purrr::pmap(
-        list(
-          torch::torch_split(out, outcome_nlevels, dim = 2),
-          torch::torch_split(batch$y, rep(1, length(outcome_nlevels)), dim = 2)
-        ),
-        ~config$loss_fn(max_constraint_output(.x, .y$squeeze(2), config$ancestor_tt))
-      )),
-      dim = 1)
-    } else {
-      # use `resolved_loss`
-      loss <- torch::torch_sum(torch::torch_stack(purrr::pmap(
-        list(
-          torch::torch_split(out, outcome_nlevels, dim = 2),
-          torch::torch_split(batch$y, rep(1, length(outcome_nlevels)), dim = 2)
-        ),
-        ~config$loss_fn(.x, .y$squeeze(2))
-      )),
-      dim = 1)
-    }
+    # use `resolved_loss`
+    loss <- torch::torch_sum(torch::torch_stack(purrr::pmap(
+      list(
+        torch::torch_split(out, outcome_nlevels, dim = 2),
+        torch::torch_split(batch$y, rep(1, length(outcome_nlevels)), dim = 2)
+      ),
+      ~config$loss_fn(.x, .y$squeeze(2))
+    )),
+    dim = 1)
+    
+  } else if (!is.null(config$ancestor)) {
+    # multi-outcome max_constraint loss ned one-hot encoding of targets
+    loss <- config$loss_fn(out, nnf_multilabel_one_hot(
+      y = batch$y,
+      outcomes = config$outcomes,
+      device = out$device
+    ))
   } else {
     if (batch$y$dtype == torch::torch_long()) {
       # classifier needs a squeeze for bce loss
@@ -513,7 +493,10 @@ tabnet_train_supervised <- function(obj, x, y, config = tabnet_config(), epoch_s
 
   # provide ancestor to torch tensor in case of hierarchical classification
   if (!is.null(config$ancestor)) {
-    config$ancestor_tt <- torch::torch_tensor(config$ancestor)$to(torch::torch_bool(), device = device)
+    if (!inherits(config$ancestor, "torch_tensor")) {
+      # config is expected to carry the tensor
+      runtime_error("ancestor was configured. Expecting a tensor but got {.cls {class(config$ancestor)}}")
+    }
   }
 
   # instantiate optimizer
@@ -579,9 +562,9 @@ tabnet_train_supervised <- function(obj, x, y, config = tabnet_config(), epoch_s
       metrics[[epoch]][["valid"]] <- transpose_metrics(valid_metrics)$loss
     }
 
-    if (config$verbose & !has_valid)
+    if (config$verbose && !has_valid)
       message(gettextf("[Epoch %03d] Loss: %3f", epoch, mean(metrics[[epoch]]$train)))
-    if (config$verbose & has_valid)
+    if (config$verbose && has_valid)
       message(gettextf("[Epoch %03d] Loss: %3f, Valid loss: %3f", epoch, mean(metrics[[epoch]]$train), mean(metrics[[epoch]]$valid)))
 
 
@@ -690,7 +673,7 @@ predict_impl_numeric <- function(obj, x, batch_size) {
 predict_impl_numeric_multiple <- function(obj, x, batch_size) {
   p <- as.matrix(predict_impl(obj, x, batch_size))
   # TODO use a cleaner function to turn matrix into vectors
-  hardhat::spruce_numeric_multiple(!!!purrr::map(1:ncol(p), ~p[,.x]))
+  hardhat::spruce_numeric_multiple(!!!purrr::map(seq_len(ncol(p)), ~p[,.x]))
 }
 
 #' single-outcome level blueprint
diff --git a/R/parsnip.R b/R/parsnip.R
index 00627a59..33d16fd6 100644
--- a/R/parsnip.R
+++ b/R/parsnip.R
@@ -539,7 +539,7 @@ multi_predict._tabnet_fit <- function(object, new_data, type = NULL, epochs = NU
     pred <- predict(object$fit, new_data, type = type, epoch = epoch)
     nms <- names(pred)
     pred[["epochs"]] <- epoch
-    pred[[".row"]] <- 1:nrow(new_data)
+    pred[[".row"]] <- seq_len(nrow(new_data))
     pred[, c(".row", "epochs", nms)]
   })
 
diff --git a/R/plot.R b/R/plot.R
index f84ea638..4d3985d5 100644
--- a/R/plot.R
+++ b/R/plot.R
@@ -41,7 +41,7 @@ autoplot.tabnet_fit <- function(object, ...) {
 
   if ("checkpoint" %in% names(collect_metrics)) {
     checkpoints <- collect_metrics %>%
-      dplyr::filter(checkpoint == TRUE, dataset == "train") %>%
+      dplyr::filter(checkpoint, dataset == "train") %>%
       dplyr::select(-checkpoint) %>%
       dplyr::mutate(size = 2)
     p +
diff --git a/R/tab-network.R b/R/tabnet_network.R
similarity index 100%
rename from R/tab-network.R
rename to R/tabnet_network.R
diff --git a/R/utils.R b/R/utils.R
index 229c5fae..9526a145 100644
--- a/R/utils.R
+++ b/R/utils.R
@@ -74,6 +74,9 @@ check_compliant_node <- function(node) {
                 Please change those names as they will lead to unexpected tabnet behavior.")
   }
   
+
+  
+  
   invisible(node)
 }
 
diff --git a/man/build_ancestor_matrix_from_outcomes.Rd b/man/build_ancestor_matrix_from_outcomes.Rd
new file mode 100644
index 00000000..96ab70ac
--- /dev/null
+++ b/man/build_ancestor_matrix_from_outcomes.Rd
@@ -0,0 +1,23 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/hardhat.R
+\name{build_ancestor_matrix_from_outcomes}
+\alias{build_ancestor_matrix_from_outcomes}
+\title{Build ancestor matrix aligned with observed outcome classes}
+\usage{
+build_ancestor_matrix_from_outcomes(x, outcomes, device = "cpu")
+}
+\arguments{
+\item{x}{A \code{data.tree::Node} object.}
+
+\item{outcomes}{A tibble with factor columns (one per hierarchy level),
+as returned by \code{hardhat::mold()$outcomes}.}
+
+\item{device}{Torch device ("cpu" or "cuda").}
+}
+\value{
+A \code{torch_tensor} of shape \verb{(1, n_classes, n_classes)}.
+}
+\description{
+Extracts class names from the outcome tibble (factor levels) and builds
+the ancestor matrix only for classes that actually appear in the data.
+}
diff --git a/man/get_constr_output.Rd b/man/get_constr_output.Rd
new file mode 100644
index 00000000..92941eb2
--- /dev/null
+++ b/man/get_constr_output.Rd
@@ -0,0 +1,22 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/loss.R
+\name{get_constr_output}
+\alias{get_constr_output}
+\title{Apply hierarchy constraints via max-pooling over descendants (MCM)}
+\usage{
+get_constr_output(x, R)
+}
+\arguments{
+\item{x}{A \code{torch_tensor} of shape \verb{(batch_size, n_classes)}.}
+
+\item{R}{A \code{torch_tensor} of shape \verb{(1, n_classes, n_classes)} where
+\code{R[1, i, j] = 1} iff class \code{i} is a descendant of class \code{j}.}
+}
+\value{
+A \code{torch_tensor} of shape \verb{(batch_size, n_classes)} with constrained outputs.
+}
+\description{
+Given neural network outputs x and ancestor matrix R, enforces that
+if a class is predicted positive, all its ancestors must also be positive.
+Implements: \verb{final_out[i] = max\{x[j] : R[i,j] = 1\}}
+}
diff --git a/man/nn_mc_loss.Rd b/man/nn_mc_loss.Rd
new file mode 100644
index 00000000..37a5ab02
--- /dev/null
+++ b/man/nn_mc_loss.Rd
@@ -0,0 +1,56 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/loss.R
+\name{nn_mc_loss}
+\alias{nn_mc_loss}
+\title{Max-Constraint Margin Loss (module)}
+\usage{
+nn_mc_loss(
+  R,
+  to_eval = NULL,
+  criterion = torch::nnf_binary_cross_entropy_with_logits,
+  reduction = "mean"
+)
+}
+\arguments{
+\item{R}{Ancestor matrix tensor of shape \verb{(1, n_classes, n_classes)}.}
+
+\item{to_eval}{Optional logical tensor of shape \verb{(n_classes,)} indicating
+which classes to include in loss computation.}
+
+\item{criterion}{Loss function module or functional to apply after constraint
+propagation. Default: \code{nn_binary_cross_entropy_with_logits()}.}
+
+\item{reduction}{(string, optional): Reduction method: \code{'none'} | \code{'mean'} | \code{'sum'}.}
+}
+\description{
+Module wrapper for \code{\link[=nnf_mc_loss]{nnf_mc_loss()}} with configurable parameters.
+Stores the ancestor matrix R and evaluation mask for reuse across batches.
+}
+\section{Shape}{
+
+\itemize{
+\item Input \code{output}: \eqn{(N, C)} where N = batch size, C = number of classes
+\item Input \code{target}: \eqn{(N, C)}, same shape as output, binary values
+\item Output: scalar by default. If \code{reduction = "none"}, then \eqn{(N, C')}
+where C' is the number of evaluated classes
+}
+}
+
+\examples{
+\dontrun{
+# Build ancestor matrix from hierarchy
+R <- build_ancestor_matrix_from_outcomes(my_tree, processed$outcomes, device = "cuda")
+
+# Create loss module
+loss_fn <- nn_mc_loss(R = R, reduction = "mean")
+
+# Forward pass
+output <- model(x)  # (batch, n_classes)
+loss <- loss_fn(output, labels)
+loss$backward()
+}
+
+}
+\seealso{
+\code{\link[=nnf_mc_loss]{nnf_mc_loss()}}, \code{\link[=build_ancestor_matrix_from_outcomes]{build_ancestor_matrix_from_outcomes()}}, \code{\link[=get_constr_output]{get_constr_output()}}
+}
diff --git a/man/nnf_mc_loss.Rd b/man/nnf_mc_loss.Rd
new file mode 100644
index 00000000..bc514e38
--- /dev/null
+++ b/man/nnf_mc_loss.Rd
@@ -0,0 +1,50 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/loss.R
+\name{nnf_mc_loss}
+\alias{nnf_mc_loss}
+\title{Max-Constraint Margin Loss (functional)}
+\usage{
+nnf_mc_loss(
+  output,
+  target,
+  R,
+  to_eval = NULL,
+  criterion = nnf_binary_cross_entropy_with_logits
+)
+}
+\arguments{
+\item{output}{A \code{torch_tensor} of raw network outputs (pre-sigmoid),
+shape \verb{(batch_size, n_classes)}.}
+
+\item{target}{Binary target labels, shape \verb{(batch_size, n_classes)}.}
+
+\item{R}{Ancestor matrix tensor of shape \verb{(1, n_classes, n_classes)} where
+\code{R[1, i, j] = 1} iff class \code{i} is a descendant of class \code{j}.}
+
+\item{to_eval}{Optional logical tensor of shape \verb{(n_classes,)} indicating
+which classes to include in the loss computation. If \code{NULL}, all classes
+are evaluated.}
+
+\item{criterion}{Loss function to apply after constraint propagation.
+Default: \code{nnf_binary_cross_entropy_with_logits} (expects raw logits).}
+}
+\value{
+A scalar \code{torch_tensor} containing the computed loss, or a tensor
+of shape \verb{(batch_size, n_classes)} if \code{reduction = "none"}.
+}
+\description{
+Computes the hierarchy-constrained loss for multi-label classification.
+Enforces that if a class is predicted positive, all its ancestors must
+also be positive, using the ancestor matrix R.
+}
+\details{
+The loss combines constrained outputs differently for positive and negative
+labels:
+\itemize{
+\item For positive labels: uses constrained output of label-weighted predictions
+\item For negative labels: uses constrained raw predictions (penalizes ancestor violations)
+}
+}
+\seealso{
+\code{\link[=nn_mc_loss]{nn_mc_loss()}}, \code{\link[=get_constr_output]{get_constr_output()}}
+}
diff --git a/man/nnf_multilabel_one_hot.Rd b/man/nnf_multilabel_one_hot.Rd
new file mode 100644
index 00000000..aceee864
--- /dev/null
+++ b/man/nnf_multilabel_one_hot.Rd
@@ -0,0 +1,23 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/loss.R
+\name{nnf_multilabel_one_hot}
+\alias{nnf_multilabel_one_hot}
+\title{Convert class_id tensor to binary one-hot tensor}
+\usage{
+nnf_multilabel_one_hot(y, outcomes, device = "cpu")
+}
+\arguments{
+\item{y}{A \code{torch_tensor} of shape \verb{(batch_size, n_levels)} containing
+1-based class indices.}
+
+\item{outcomes}{A tibble with factor columns (as from \code{hardhat::mold()$outcomes}).}
+
+\item{device}{Torch device.}
+}
+\value{
+A \code{torch_tensor} of shape \verb{(batch_size, n_classes)} with binary values.
+}
+\description{
+Transforms a tensor of class indices (one column per hierarchy level)
+into a binary tensor where each column corresponds to a class.
+}
diff --git a/man/tabnet_config.Rd b/man/tabnet_config.Rd
index d20ecd87..7295409d 100644
--- a/man/tabnet_config.Rd
+++ b/man/tabnet_config.Rd
@@ -1,5 +1,5 @@
 % Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/model.R
+% Please edit documentation in R/model_training.R
 \name{tabnet_config}
 \alias{tabnet_config}
 \title{Configuration for TabNet models}
diff --git a/man/tabnet_explain.Rd b/man/tabnet_explain.Rd
index f750c5f5..1327c039 100644
--- a/man/tabnet_explain.Rd
+++ b/man/tabnet_explain.Rd
@@ -1,5 +1,5 @@
 % Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/explain.R
+% Please edit documentation in R/model_explain.R
 \name{tabnet_explain}
 \alias{tabnet_explain}
 \alias{tabnet_explain.default}
diff --git a/man/tabnet_nn.Rd b/man/tabnet_nn.Rd
index ae5e8f76..3fa483a1 100644
--- a/man/tabnet_nn.Rd
+++ b/man/tabnet_nn.Rd
@@ -1,5 +1,5 @@
 % Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/tab-network.R
+% Please edit documentation in R/tabnet_network.R
 \name{tabnet_nn}
 \alias{tabnet_nn}
 \title{TabNet Model Architecture}
diff --git a/man/tabnet_pretrain.Rd b/man/tabnet_pretrain.Rd
index 5c1e42e7..b3777ae5 100644
--- a/man/tabnet_pretrain.Rd
+++ b/man/tabnet_pretrain.Rd
@@ -15,7 +15,7 @@ tabnet_pretrain(x, ...)
 
 \method{tabnet_pretrain}{data.frame}(
   x,
-  y,
+  y = NULL,
   tabnet_model = NULL,
   config = tabnet_config(),
   ...,
diff --git a/po/R-fr.po b/po/R-fr.po
index 14b4b8a1..efa406ad 100644
--- a/po/R-fr.po
+++ b/po/R-fr.po
@@ -225,7 +225,7 @@ msgid ""
 "                Please change those names as they will lead to unexpected "
 "tabnet behavior."
 msgstr ""
-"Les attributs ou noms de colonnes dans l’objet hiérarchique fournit utilise "
+"Les `attributs` (noms de colonne) dans l’objet hiérarchique fournit utilisent "
 "les noms réservés suivants : {.vars {actual_names[actual_names %in% "
 "reserved_names]}}. Veuillez changer ces noms pour éviter un comportement "
 "imprévisible de TabNet."
diff --git a/tests/testthat/_snaps/pretraining.md b/tests/testthat/_snaps/model_pretraining.md
similarity index 100%
rename from tests/testthat/_snaps/pretraining.md
rename to tests/testthat/_snaps/model_pretraining.md
diff --git a/tests/testthat/helper-tensor.R b/tests/testthat/helper-tensor.R
index 31b5c9bd..534bff14 100644
--- a/tests/testthat/helper-tensor.R
+++ b/tests/testthat/helper-tensor.R
@@ -38,7 +38,7 @@ expect_no_error <- function(object, ...) {
 
 expect_tensor <- function(object) {
   expect_true(torch:::is_torch_tensor(object))
-  expect_no_error(torch::as_array(object$to(device = "cpu")))
+  expect_no_error(torch::as_array(object$to_dense()$to(device = "cpu")))
 }
 
 expect_equal_to_r <- function(object, expected, ...) {
@@ -50,6 +50,13 @@ expect_tensor_shape <- function(object, expected) {
   expect_equal(object$shape, expected)
 }
 
+
+expect_tensor_dtype <- function(object, expected_dtype) {
+  expect_tensor(object)
+  expect_true(object$dtype == expected_dtype)
+}
+
+
 expect_undefined_tensor <- function(object) {
   # TODO
 }
diff --git a/tests/testthat/setup.R b/tests/testthat/setup.R
index 4b78d736..95527413 100644
--- a/tests/testthat/setup.R
+++ b/tests/testthat/setup.R
@@ -13,7 +13,7 @@ y <- ames[ids,]$Sale_Price
 
 # ames common models
 ames_pretrain <- tabnet_pretrain(x, y, epoch = 2, checkpoint_epochs = 1)
-ames_pretrain_vsplit <- tabnet_pretrain(x, y, epochs = 3, valid_split=.2,
+ames_pretrain_vsplit <- tabnet_pretrain(x, y, epochs = 3, valid_split=0.2,
                                         num_steps = 1, attention_width = 1, num_shared = 1, num_independent = 1)
 ames_fit <- tabnet_fit(x, y, epochs = 5 , checkpoint_epochs = 2)
 ames_fit_vsplit <- tabnet_fit(x, y, tabnet_model=ames_pretrain_vsplit, epochs = 3,
@@ -38,8 +38,12 @@ attr_fitted_vsplit <- tabnet_fit(attrix, attriy, epochs = 12, valid_split=0.3)
 utils::data("acme", package = "data.tree")
 acme_df <-  data.tree::ToDataFrameTypeCol(acme, acme$attributesAll) %>%
   select(-starts_with("level_"))
+# acme2 <- acme$clone()
+# acme2$RemoveAttribute("level_3")
 
 attrition_tree <- attrition %>%
+  # ensure pure tree
+  filter(!(Department %in% c("Sales", "Research_Development") & JobRole == "Manager")) %>%
   tibble::rowid_to_column() %>%
   mutate(pathString = paste("attrition", Department, JobRole, rowid, sep = "/")) %>%
   select(-Department, -JobRole, -rowid) %>%
diff --git a/tests/testthat/test-hardhat_hierarchical.R b/tests/testthat/test-hardhat_hierarchical.R
index a6fa0983..91a7d4b2 100644
--- a/tests/testthat/test-hardhat_hierarchical.R
+++ b/tests/testthat/test-hardhat_hierarchical.R
@@ -1,87 +1,4 @@
-test_that("C-HMCNN get_constr_output works ", {
-  x <- torch::torch_rand(c(2,4))
-  R <- torch::torch_tril(torch::torch_zeros(c(4,4))$bernoulli(p = 0.2) + torch::torch_diag(rep(1,4)))$to(dtype = torch::torch_bool())
-  expect_no_error(
-    constr_output <- get_constr_output(x, R)
-  )
-  expect_tensor_shape(
-    constr_output, x$shape
-  )
-  # expect_equal(
-  #   constr_output$dtype, torch_tensor(0.1)$dtype
-  # )
-
-  R <- torch::torch_zeros(c(4,4))$to(dtype = torch::torch_bool())
-  expect_equal_to_tensor(
-    get_constr_output(x, R), torch::torch_zeros_like(x)
-  )
-})
-
-test_that("C-HMCNN max_constraint_output works ", {
-  output <- torch::torch_rand(c(3, 5))
-  labels <- torch::torch_diag(rep(1,5))[1:3, ]$to(dtype = torch::torch_bool())
-  ancestor <- torch::torch_tril(torch::torch_zeros(c(5, 5))$bernoulli(p = 0.2) )$to(dtype = torch::torch_bool())
-
-  expect_no_error(
-    MC_output <- max_constraint_output(output, labels, ancestor)
-  )
-  expect_tensor_shape(
-    MC_output, output$shape
-  )
-  # max_constraint_output is not identity
-  expect_not_equal_to_tensor(
-    MC_output, output
-  )
-  # max_constraint_output provides more than 35% null values
-  expect_gte(
-    as.matrix(torch::torch_sum(MC_output == 0), device="cpu"), .30 * output$shape[1] * output$shape[2]
-  )
-})
-
-test_that("node_to_df works ", {
-  expect_no_error(
-    node_to_df(acme)
-  )
-  expect_no_error(
-    attrition_df <- node_to_df(attrition_tree)
-  )
-  # node_to_df removes first and last level of the hierarchy
-  outcome_levels <- paste0("level_", seq(2, attrition_tree$height - 1))
-  expect_equal(names(attrition_df$y), outcome_levels)
-
-  # node_to_df do not shuffle outcome rows
-  df <- tibble(pred_1 = seq(1,26), pred_2 = seq(26,1),
-               level_2 = factor(LETTERS[1:26]), level_3 = factor(letters[26:1]))
-  df_node_df <- df %>%
-    mutate(pathString = paste("synth", level_2, level_3, level_3, sep = "/")) %>%
-    select(-level_2, -level_3) %>%
-    as.Node() %>%
-    node_to_df()
-
-  expect_equal(df_node_df$y %>% as_tibble(), df %>% select(starts_with("level_")))
-  expect_equal(df_node_df$x %>% as_tibble(), df %>% select(starts_with("pred_")))
-
-})
-
-
-test_that("Training hierarchical classification for {data.tree} Node", {
-
-  expect_no_error(
-    fit <- tabnet_fit(acme, epochs = 1)
-  )
-  expect_no_error(
-    result <- predict(fit, acme_df, type = "prob")
-  )
-
-  expect_equal(ncol(result), 3)
-  outcome_levels <-levels(fit$blueprint$ptypes$outcomes[[1]])
-  # we get back outcomes vars with a `.pred_` prefix
-  expect_equal(stringr::str_remove(names(result), ".pred_"), outcome_levels)
-  expect_no_error(
-    result <- predict(fit, acme_df)
-  )
-  expect_equal(ncol(result), 1)
-
+test_that("Training hierarchical classification for {data.tree} Node attrition_tree", {
   expect_no_error(
     fit <- tabnet_fit(attrition_tree, epochs = 1)
   )
@@ -91,7 +8,7 @@ test_that("Training hierarchical classification for {data.tree} Node", {
 
   expect_equal(ncol(result), 2) # 2 outcomes levels_
 
-  outcome_nlevels <- purrr::map_dbl(fit$blueprint$ptypes$outcomes, ~length(levels(.x)))
+  outcome_nlevels <- purrr::map_dbl(fit$blueprint$ptypes$outcomes, ~nlevels(.x))
   # we get back outcomes vars with a `.pred_` prefix
   expect_equal(stringr::str_remove(names(result), ".pred_"), names(outcome_nlevels))
 
@@ -106,14 +23,15 @@ test_that("Training hierarchical classification for {data.tree} Node with valida
   expect_no_error(
     fit <- tabnet_fit(attrition_tree, valid_split = 0.2, epochs = 1)
   )
-
+  expect_true( "ancestor" %in% names(fit$fit$config))
+  
   expect_no_error(
     result <- predict(fit, attrition_tree, type = "prob")
   )
 
   expect_equal(ncol(result), 2) # 2 outcomes levels_
 
-  outcome_nlevels <- purrr::map_dbl(fit$blueprint$ptypes$outcomes, ~length(levels(.x)))
+  outcome_nlevels <- purrr::map_dbl(fit$blueprint$ptypes$outcomes, ~nlevels(.x))
   # we get back outcomes vars with a `.pred_` prefix
   expect_equal(stringr::str_remove(names(result), ".pred_"), names(outcome_nlevels))
 
@@ -144,11 +62,25 @@ test_that("hierarchical classification for {data.tree} Node is explainable", {
 
 })
 
+test_that("Training hierarchical classification for {data.tree} Node for starwars", {
+
+  starwars_tree <- starwars %>% 
+    rename(`_name` = "name", `_height` = "height") %>% 
+    mutate(species = coalesce(species, "Unknown_Species"),
+           sex     = coalesce(sex, "Unknown_Sex"),
+           pathString = paste("StarWars_characters", species, sex, `_name`, sep = "/")) %>%
+    as.Node()
+  
+  expect_no_error(  fit <- tabnet_fit(starwars_tree, epochs = 1))
+})
+
 test_that("we properly check non-compliant colnames", {
 
   # try to use starwars dataset with two forbidden column name
   starwars_tree <- starwars %>%
-    mutate(pathString = paste("tree", species, homeworld, `name`, sep = "/"))
+    mutate(species = coalesce(species, "Unknown_Species"),
+           sex     = coalesce(sex, "Unknown_Sex"),
+           pathString = paste("tree", species, homeworld, `name`, sep = "/"))
   expect_error(
     check_compliant_node(starwars_tree)
     ,"reserved names")
diff --git a/tests/testthat/test-hardhat_multi-outcome.R b/tests/testthat/test-hardhat_multi-outcome.R
index fa3eafd2..d1a79d43 100644
--- a/tests/testthat/test-hardhat_multi-outcome.R
+++ b/tests/testthat/test-hardhat_multi-outcome.R
@@ -54,7 +54,7 @@ test_that("Training multilabel classification from data.frame", {
   )
 
   expect_equal(ncol(result), 3)
-  outcome_nlevels <- purrr::map_dbl(fit$blueprint$ptypes$outcomes, ~length(levels(.x)))
+  outcome_nlevels <- purrr::map_dbl(fit$blueprint$ptypes$outcomes, ~nlevels(.x))
   # we get back outcomes vars with a `.pred_` prefix
   expect_equal(stringr::str_remove(names(result), ".pred_"), names(outcome_nlevels))
 
@@ -82,7 +82,7 @@ test_that("Training multilabel classification from formula", {
   )
 
   expect_equal(ncol(result), 2)
-  outcome_nlevels <- purrr::map_dbl(fit$blueprint$ptypes$outcomes, ~length(levels(.x)))
+  outcome_nlevels <- purrr::map_dbl(fit$blueprint$ptypes$outcomes, ~nlevels(.x))
   # we get back outcomes vars with a `.pred_` prefix
   expect_equal(stringr::str_remove(names(result), ".pred_"), names(outcome_nlevels))
 
@@ -108,7 +108,7 @@ test_that("Training multilabel classification from recipe", {
   )
   expect_equal(ncol(result), 2)
 
-  outcome_nlevels <- purrr::map_dbl(fit$blueprint$ptypes$outcomes, ~length(levels(.x)))
+  outcome_nlevels <- purrr::map_dbl(fit$blueprint$ptypes$outcomes, ~nlevels(.x))
   expect_equal(stringr::str_remove(names(result), ".pred_class_"), names(outcome_nlevels))
 
 })
@@ -126,7 +126,7 @@ test_that("Training multilabel classification from data.frame with validation sp
 
   expect_equal(ncol(result), 3)
 
-  outcome_nlevels <- purrr::map_dbl(fit$blueprint$ptypes$outcomes, ~length(levels(.x)))
+  outcome_nlevels <- purrr::map_dbl(fit$blueprint$ptypes$outcomes, ~nlevels(.x))
   # we get back outcomes vars with a `.pred_` prefix
   expect_equal(stringr::str_remove(names(result), ".pred_"), names(outcome_nlevels))
 
diff --git a/tests/testthat/test-hierarchical_utils.R b/tests/testthat/test-hierarchical_utils.R
new file mode 100644
index 00000000..b84ba267
--- /dev/null
+++ b/tests/testthat/test-hierarchical_utils.R
@@ -0,0 +1,176 @@
+test_that("returns correct shape and type for a simple 2-level hierarchy", {
+  # Arbre : Root -> {A, B}, A -> {C1, C2}, B -> {D1, D2}
+  tree_df <- data.frame(pathString = c(
+    "Root/A/C1", "Root/A/C2", "Root/B/D1", "Root/B/D2"
+  ))
+  tree <- as.Node(tree_df)
+  
+  outcomes <- tibble::tibble(
+    level_2 = factor(c("A", "A", "B", "B")),
+    level_3 = factor(c("C1", "C2", "D1", "D2"))
+  )
+  
+  result <- build_ancestor_matrix_from_outcomes(tree, outcomes)
+  
+  # 2 level_2 classes + 4 level_3 = 6 classes
+  expect_tensor(result)
+  expect_tensor_shape(result, c(1, 6, 6))
+  expect_tensor_dtype(result, torch::torch_double())
+})
+
+test_that("handles ghost classes (present in tree but absent from outcomes)", {
+  # Tree with a "C" branch not in the outcomes
+  tree_df <- data.frame(pathString = c(
+    "Root/A/C1", "Root/A/C2", 
+    "Root/B/D1", "Root/B/D2", 
+    "Root/C/E1", "Root/C/E2"
+  ))
+  tree <- as.Node(tree_df)
+  
+  # Outcomes shall only contain A and B (C is a "ghost class")
+  outcomes <- tibble::tibble(
+    level_2 = factor(c("A", "A", "B", "B")), 
+    level_3 = factor(c("C1", "C2", "D1", "D2"))
+  )
+  
+  result <- build_ancestor_matrix_from_outcomes(tree, outcomes)
+  expect_tensor_shape(result, c(1, 6, 6))
+  
+  # Check of transitive loop : A(1) is ancestor of C1(3) and C2(4)
+  # Order : A(1), B(2), C1(3), C2(4), D1(5), D2(6)
+  expect_equal_to_r(result[1, 3, 1], 1) 
+  expect_equal_to_r(result[1, 4, 1], 1)
+  expect_equal_to_r(result[1, 5, 2], 1) 
+  expect_equal_to_r(result[1, 6, 2], 1)
+  expect_equal_to_r(result[1, 3, 3], 1) # Self-loop
+})
+
+test_that("handles non-unique names across different hierarchy levels", {
+  # "Manager" may exist at level_2 (Department) and level_3 (JobRole)
+  tree_df <- data.frame(pathString = c(
+    "Root/Manager/Rep", 
+    "Root/IT/Manager"
+  ))
+  tree <- as.Node(tree_df)
+  
+  outcomes <- tibble::tibble(
+    level_2 = factor(c("Manager", "IT")),
+    level_3 = factor(c("Rep", "Manager"))
+  )
+  
+  result <- build_ancestor_matrix_from_outcomes(tree, outcomes)
+  
+  # 2 + 2 = 4 classes
+  expect_tensor_shape(result, c(1L, 4L, 4L))
+
+  # Order : Manager_lvl2(1), IT(2), Rep(3), Manager_lvl3(4)
+  # Manager_lvl2(1) is ancestor of Rep(3)
+  expect_equal_to_r(result[1, 3, 1], 1)
+  # IT(2) is ancestor of Manager_lvl3(4)
+  expect_equal_to_r(result[1, 4, 2], 1)
+})
+
+test_that("throws an explicit error when a factor level is missing from the tree", {
+  tree_df <- data.frame(pathString = c("Root/A/C1", "Root/A/C2"))
+  tree <- as.Node(tree_df)
+  
+  outcomes <- tibble::tibble(
+    level_2 = factor(c("A", "X")), # "X" is not in the tree
+    level_3 = factor(c("C1", "C2"))
+  )
+  
+  expect_error(
+    build_ancestor_matrix_from_outcomes(tree, outcomes),
+    "not found"
+  )
+})
+
+test_that("throws an error when outcomes contains no factor levels", {
+  tree_df <- data.frame(pathString = c("Root/A/C1"))
+  tree <- as.Node(tree_df)
+  
+  outcomes <- tibble::tibble(
+    level_2 = factor(character(0)),
+    level_3 = factor(character(0))
+  )
+  
+  expect_error(
+    build_ancestor_matrix_from_outcomes(tree, outcomes),
+    "No factor levels"
+  )
+})
+
+test_that("preserves the exact class order defined in outcomes factors", {
+  # Dans l'arbre, B est défini avant A
+  tree_df <- data.frame(pathString = c(
+    "Root/B/D1", "Root/B/D2", 
+    "Root/A/C1", "Root/A/C2"
+  ))
+  tree <- as.Node(tree_df)
+  
+  # Mais dans outcomes, A est explicitement avant B
+  outcomes <- tibble::tibble(
+    level_2 = factor(c("A", "A", "B", "B"), levels = c("A", "B")),
+    level_3 = factor(c("C1", "C2", "D1", "D2"), levels = c("C1", "C2", "D1", "D2"))
+  )
+  
+  result <- build_ancestor_matrix_from_outcomes(tree, outcomes)
+
+  # Order is given by outcomes : A(1), B(2), C1(3), C2(4), D1(5), D2(6)
+  # A(1) is ancestor of C1(3) et C2(4)
+  expect_equal_to_r(result[1, 3, 1], 1)
+  expect_equal_to_r(result[1, 4, 1], 1)
+  
+  # B(2) is ancestor of D1(5) et D2(6)
+  expect_equal_to_r(result[1, 5, 2], 1)
+  expect_equal_to_r(result[1, 6, 2], 1)
+})
+
+test_that("handles a flat single-level hierarchy correctly", {
+  tree_df <- data.frame(pathString = c("Root/A", "Root/B", "Root/C"))
+  tree <- as.Node(tree_df)
+  
+  outcomes <- tibble::tibble(
+    level_2 = factor(c("A", "B", "C"))
+  )
+  
+  result <- build_ancestor_matrix_from_outcomes(tree, outcomes)
+  expect_equal(result$shape, c(1L, 3L, 3L))
+  
+  # no hierarchy, only self-loops
+  expect_equal_to_r(result$squeeze(1), diag(3))
+})
+
+test_that("computes full transitive closure for deep hierarchies (3+ levels)", {
+  tree_df <- data.frame(pathString = c(
+    "Root/L1_A/L2_A1/L3_A1a", 
+    "Root/L1_A/L2_A1/L3_A1b", 
+    "Root/L1_A/L2_A2/L3_A2a", 
+    "Root/L1_B/L2_B1/L3_B1a"
+  ))
+  tree <- as.Node(tree_df)
+  
+  outcomes <- tibble::tibble(
+    level_2 = factor(c("L1_A", "L1_A", "L1_A", "L1_B")),
+    level_3 = factor(c("L2_A1","L2_A1", "L2_A2", "L2_B1")),
+    level_4 = factor(c("L3_A1a", "L3_A1b", "L3_A2a", "L3_B1a"))
+  )
+  
+  result <- build_ancestor_matrix_from_outcomes(tree, outcomes)
+  expect_equal(result$shape, c(1L, 9L, 9L))
+  
+  # Order : L1_A(1), L1_B(2), L2_A1(3), L2_A2(4), L2_B1(5), 
+  #         L3_A1a(6), L3_A1b(7), L3_A2a(8), L3_B1a(9)
+  
+  # L1_A(1) is a transitive ancestor of all the sub-tree A
+  expect_equal_to_r(result[1, 3, 1], 1) # -> L2_A1
+  expect_equal_to_r(result[1, 6, 1], 1) # -> L3_A1a
+  expect_equal_to_r(result[1, 8, 1], 1) # -> L3_A2a
+  
+  # L2_A1(3) is an ancestor of all its direct children
+  expect_equal_to_r(result[1, 6, 3], 1) # -> L3_A1a
+  expect_equal_to_r(result[1, 7, 3], 1) # -> L3_A1b
+  
+  # Self-loops on the diagonal (substracting the eye don't go to negative values)
+  expect_true((result$squeeze() - torch::torch_eye(9))$min()$item() >= 0)
+})
\ No newline at end of file
diff --git a/tests/testthat/test-loss.R b/tests/testthat/test-loss.R
index cff90a32..4fb7db5b 100644
--- a/tests/testthat/test-loss.R
+++ b/tests/testthat/test-loss.R
@@ -62,12 +62,207 @@ test_that("nn_aum_loss works as expected with {n, 2} shape prediction", {
   output <- aum_loss(pred_tensor, label_tensor)
   output$backward()
   
-  
   expect_tensor(output)
   expect_equal_to_r(output >= 0, TRUE) 
   expect_false(rlang::is_null(output$grad_fn))
   expect_equal(output$dim(), 0)
+})
+
+
+test_that("get_constr_output handles basic 2D input with identity constraint", {
+  m <- matrix(c(1, 2, 
+                3, 4), nrow = 2, ncol = 2)
+  x <- torch_tensor(m, dtype = torch::torch_float32())
+  R <- torch::torch_eye(2)
+  result <- get_constr_output(x, R)
+  expect_tensor(result)
+  expect_tensor_shape(result, c(2, 2))
+  expect_equal_to_r(result, m)
+})
+
+test_that("get_constr_output applies hierarchy constraint correctly", {
+  x <- torch_tensor(matrix(c(1, 5, 
+                             3, 2), nrow = 2, ncol = 2, byrow = TRUE), dtype = torch::torch_float64())
+  R <- torch_tensor(matrix(c(1, 1, 
+                             0, 1), nrow = 2, ncol = 2, byrow = TRUE), dtype = torch::torch_float64())
+  result <-get_constr_output(x, R)
+  expect_tensor_shape(result, c(2, 2))
+  expected <- matrix(c(5, 5, 3, 2), nrow = 2, ncol = 2, byrow = TRUE)
+  expect_equal_to_r(result, expected, tolerance = 1e-6)
+})
+
+test_that("get_constr_output preserves input dtype", {
+  x_f32 <- torch_tensor(matrix(1:4, nrow = 2), dtype = torch::torch_float32())
+  x_f64 <- torch_tensor(matrix(1:4, nrow = 2), dtype = torch::torch_float64())
+  R <- torch::torch_eye(2)
+  expect_tensor_dtype(get_constr_output(x_f32, R), torch::torch_float64())
+  expect_tensor_dtype(get_constr_output(x_f64, R), torch::torch_float64())
+})
+
+test_that("get_constr_output handles batch dimension correctly", {
+  x <- torch_tensor(matrix(1:12, nrow = 3, ncol = 4))
+  R <- torch_tensor(matrix(c(1, 1, 0, 0, 
+                             1, 1, 0, 0, 
+                             0, 0, 1, 1, 
+                             0, 0, 1, 1), nrow = 4, ncol = 4, byrow = TRUE))
+  result <-get_constr_output(x, R)
+  expect_tensor_shape(result, c(3, 4))
   
+  for (i in 1:3) {
+    row_result <- as_array(result[i, ])
+    max_grp1 <- max(as_array(x[i, 1:2]))
+    max_grp2 <- max(as_array(x[i, 3:4]))
+    
+    expect_equal(row_result[1:2], rep(max_grp1, 2), tolerance = 1e-6)
+    expect_equal(row_result[3:4], rep(max_grp2, 2), tolerance = 1e-6)
+  }
+})
+test_that("get_constr_output works with single sample", {
+  x <- torch_tensor(matrix(c(2, 1, 4, 3), nrow = 1, ncol = 4, byrow = TRUE))
+  R <- torch_tensor(matrix(c(1, 1, 0, 0, 
+                             1, 1, 0, 0, 
+                             0, 0, 1, 1, 
+                             0, 0, 1, 1), nrow = 4, ncol = 4, byrow = TRUE))
+  result <-get_constr_output(x, R)
+  expect_tensor_shape(result, c(1, 4))
+  expected <- matrix(c(2, 2, 4, 4), nrow = 1, byrow = TRUE)
+  expect_equal_to_r(result, expected)
 })
 
+test_that("get_constr_output handles all-zeros constraint matrix", {
+  x <- torch_tensor(matrix(1:6, nrow = 2, ncol = 3))
+  R <- torch::torch_zeros(c(3, 3))
+  result <-get_constr_output(x, R)
+  expect_tensor_shape(result, c(2, 3))
+  expect_equal_to_r(result, matrix(0, nrow = 2, ncol = 3))
+})
 
+test_that("get_constr_output handles all-ones constraint matrix", {
+  x <- torch_tensor(matrix(c(1, 5, 3, 
+                             2, 4, 6), nrow = 2, ncol = 3, byrow = TRUE))
+  R <- torch::torch_ones(c(3, 3))
+  result <-get_constr_output(x, R)
+  expect_tensor_shape(result, c(2, 3))
+  # Each row is filled with its own row-wise maximum
+  expected <- matrix(c(5, 5, 5, 
+                       6, 6, 6), nrow = 2, ncol = 3, byrow = TRUE)
+  expect_equal_to_r(result, expected, tolerance = 1e-6)
+})
+
+test_that("get_constr_output throws error for dimension mismatch", {
+  x <- torch_tensor(matrix(1:4, nrow = 2, ncol = 2))
+  R <- torch::torch_eye(3)
+  expect_error(get_constr_output(x, R), "must match the existing size")
+})
+
+test_that("get_constr_output throws error for non-2D R", {
+  x <- torch_tensor(matrix(1:4, nrow = 2, ncol = 2))
+  R <- torch_tensor(array(1:8, dim = c(1, 2, 2, 2)))
+  expect_error(get_constr_output(x, R), "dimension")
+})
+
+test_that("get_constr_output handles negative values correctly", {
+  x <- torch_tensor(matrix(c(-5, -1, 
+                             -3, -2), nrow = 2, ncol = 2, byrow = TRUE))
+  R <- torch_tensor(matrix(c(1, 1, 
+                             0, 1), nrow = 2, ncol = 2, byrow = TRUE))
+  result <- get_constr_output(x, R)
+  expected <- matrix(c(-1, 0, 
+                       -2, 0), nrow = 2, ncol = 2, byrow = TRUE)
+  expect_equal_to_r(result, expected)
+})
+
+test_that("nn_mc_loss resolves functional criterion at initialization", {
+  R <- torch::torch_eye(3)$unsqueeze(1)
+  
+  # Functional criterion
+  loss_fn <- nn_mc_loss(
+    R = R,
+    criterion = torch::nnf_binary_cross_entropy_with_logits,
+    reduction = "mean"
+  )
+  
+  expect_true(rlang::is_function(loss_fn$criterion_fn))
+  
+  output <- torch::torch_randn(2, 3, requires_grad = TRUE)
+  target <- torch::torch_randint(0, 2, c(2, 3))$to(dtype = torch::torch_double())
+  
+  expect_no_error(loss <- loss_fn(output, target))
+  expect_tensor(loss)
+})
+
+test_that("nn_mc_loss resolves nn_module criterion at initialization (default)", {
+  R <- torch::torch_eye(3)$unsqueeze(1)
+  
+  # Functional criterion
+  loss_fn <- nn_mc_loss(R = R)
+  
+  expect_true(rlang::is_function(loss_fn$criterion_fn))
+  
+  output <- torch::torch_randn(2, 3, requires_grad = TRUE)
+  target <- torch::torch_randint(0, 2, c(2, 3))$to(dtype = torch::torch_double())
+  
+  expect_no_error(loss <- loss_fn(output, target))
+  expect_tensor(loss)
+})
+
+test_that("nn_mc_loss can use already instanciated nn_module criterion", {
+  R <- torch::torch_eye(3)$unsqueeze(1)
+  
+  # Module criterion
+  loss_fn <- nn_mc_loss(
+    R = R,
+    criterion = torch::nn_bce_with_logits_loss(),
+    reduction = "mean"
+  )
+  
+  expect_true(rlang::is_function(loss_fn$criterion_fn))
+  
+  output <- torch::torch_randn(2, 3, requires_grad = TRUE)
+  target <- torch::torch_randint(0, 2, c(2, 3))$to(dtype = torch::torch_double())
+  
+  expect_no_error(loss <- loss_fn(output, target))
+  expect_tensor(loss)
+})
+
+test_that("nn_mc_loss errors on invalid criterion type", {
+  R <- torch::torch_eye(3)$unsqueeze(1)$to(dtype = torch::torch_double())
+  
+  expect_error(
+    nn_mc_loss(R = R, criterion = "not_a_valid_criterion"),
+    "must be a function or an"
+  )
+})
+
+test_that("nn_mc_loss warns on reduction mismatch for module criterion", {
+  R <- torch::torch_eye(3)$unsqueeze(1)$to(dtype = torch::torch_double())
+  
+  # Module with 'sum' reduction, but loss asks for 'mean'
+  expect_warning(
+    nn_mc_loss(
+      R = R,
+      criterion = torch::nn_bce_with_logits_loss(reduction = "sum"),
+      reduction = "mean"
+    ),
+    "The criterion module has reduction"
+  )
+})
+
+test_that("nn_mc_loss backward pass works without inplace errors", {
+  R <- torch::torch_eye(3)$unsqueeze(1)$to(dtype = torch::torch_double())
+  
+  loss_fn <- nn_mc_loss(R = R, reduction = "mean")
+  
+  output <- torch::torch_randn(2, 3, requires_grad = TRUE)
+  target <- torch::torch_randint(0, 2, c(2, 3))$to(dtype = torch::torch_double())
+  
+  # Forward
+  loss <- loss_fn(output, target)
+  
+  # Backward should not throw inplace error
+  expect_no_error(loss$backward())
+  
+  # Gradients should be computed
+  expect_true(!is.null(output$grad))
+  expect_tensor_shape(output$grad, output$shape)
+})
\ No newline at end of file
diff --git a/tests/testthat/test-explain.R b/tests/testthat/test-model_explain.R
similarity index 98%
rename from tests/testthat/test-explain.R
rename to tests/testthat/test-model_explain.R
index 549afed4..db0ace3b 100644
--- a/tests/testthat/test-explain.R
+++ b/tests/testthat/test-model_explain.R
@@ -52,7 +52,7 @@ test_that("explain works for dataframe, formula and recipe", {
 
 
   # formula
-  tabnet_pretrain <- tabnet_pretrain(Sale_Price ~., data=small_ames, epochs = 3, valid_split=.2,
+  tabnet_pretrain <- tabnet_pretrain(Sale_Price ~., data=small_ames, epochs = 3, valid_split=0.2,
                                      num_steps = 1, attention_width = 1, num_shared = 1, num_independent = 1)
   expect_no_error(
     tabnet_explain(tabnet_pretrain, new_data=small_ames)
@@ -69,7 +69,7 @@ test_that("explain works for dataframe, formula and recipe", {
     step_zv(all_predictors()) %>%
     step_normalize(all_numeric_predictors())
 
-  tabnet_pretrain <- tabnet_pretrain(rec, data=small_ames, epochs = 3, valid_split=.2,
+  tabnet_pretrain <- tabnet_pretrain(rec, data=small_ames, epochs = 3, valid_split=0.2,
                                      num_steps = 1, attention_width = 1, num_shared = 1, num_independent = 1)
   expect_no_error(
     tabnet_explain(tabnet_pretrain, new_data=small_ames)
diff --git a/tests/testthat/test-pretraining.R b/tests/testthat/test-model_pretraining.R
similarity index 100%
rename from tests/testthat/test-pretraining.R
rename to tests/testthat/test-model_pretraining.R
diff --git a/tests/testthat/test-model.R b/tests/testthat/test-model_training.R
similarity index 100%
rename from tests/testthat/test-model.R
rename to tests/testthat/test-model_training.R
diff --git a/tests/testthat/test-parsnip.R b/tests/testthat/test-parsnip.R
index 742281e8..f322bcbe 100644
--- a/tests/testthat/test-parsnip.R
+++ b/tests/testthat/test-parsnip.R
@@ -98,6 +98,7 @@ test_that("Check we can finalize a workflow from a tune_grid", {
 
   model <- tabnet(epochs = tune(), checkpoint_epochs = 1) %>%
     parsnip::set_mode("regression") %>%
+    parsnip::set_args(epochs = 2) %>%
     parsnip::set_engine("torch")
 
   wf <- workflows::workflow() %>%
@@ -134,7 +135,7 @@ test_that("tabnet grid reduction - torch", {
 
   expect_equal(reg_grid_smol$epochs, rep(3, 2))
   expect_equal(reg_grid_smol$penalty, 1:2)
-  for (i in 1:nrow(reg_grid_smol)) {
+  for (i in seq_len(nrow(reg_grid_smol))) {
     expect_equal(reg_grid_smol$.submodels[[i]], list(epochs = 1:2))
   }
 
@@ -155,7 +156,7 @@ test_that("tabnet grid reduction - torch", {
   expect_equal(reg_grid_extra_smol$epochs, rep(3, 6))
   expect_equal(reg_grid_extra_smol$penalty, rep(1:2, each = 3))
   expect_equal(reg_grid_extra_smol$batch_size, rep(10:12, 2))
-  for (i in 1:nrow(reg_grid_extra_smol)) {
+  for (i in seq_len(nrow(reg_grid_extra_smol))) {
     expect_equal(reg_grid_extra_smol$.submodels[[i]], list(epochs = 1:2))
   }
 
@@ -172,7 +173,7 @@ test_that("tabnet grid reduction - torch", {
 
   expect_equal(no_sub_smol$epochs, rep(1, 2))
   expect_equal(no_sub_smol$penalty, 1:2)
-  for (i in 1:nrow(no_sub_smol)) {
+  for (i in seq_len(nrow(no_sub_smol))) {
     expect_length(no_sub_smol$.submodels[[i]], 0)
   }
 
@@ -184,7 +185,7 @@ test_that("tabnet grid reduction - torch", {
 
   expect_equal(reg_grid_smol$Amos, rep(3, 2))
   expect_equal(reg_grid_smol$penalty, 1:2)
-  for (i in 1:nrow(reg_grid_smol)) {
+  for (i in seq_len(nrow(reg_grid_smol))) {
     expect_equal(reg_grid_smol$.submodels[[i]], list(Amos = 1:2))
   }
 
@@ -202,7 +203,7 @@ test_that("tabnet grid reduction - torch", {
   expect_equal(reg_grid_smol$`Ade Tukunbo`, rep(3, 4))
   expect_equal(reg_grid_smol$penalty, rep(1:2, each = 2))
   expect_equal(reg_grid_smol$` \t123`, rep(10:11, 2))
-  for (i in 1:nrow(reg_grid_smol)) {
+  for (i in seq_len(nrow(reg_grid_smol))) {
     expect_equal(reg_grid_smol$.submodels[[i]], list(`Ade Tukunbo` = 1:2))
   }
 })
diff --git a/tests/testthat/test_translations.R b/tests/testthat/test-translations.R
similarity index 100%
rename from tests/testthat/test_translations.R
rename to tests/testthat/test-translations.R
diff --git a/vignettes/Hierarchical_classification.Rmd b/vignettes/Hierarchical_classification.Rmd
index ab7bee93..b2af4229 100644
--- a/vignettes/Hierarchical_classification.Rmd
+++ b/vignettes/Hierarchical_classification.Rmd
@@ -25,20 +25,32 @@ library(tibble)
 set.seed(202307)
 ```
 
-## Data preparation
+## Data format
 
 The supported data format for hierarchical classification is the `Node` object format from package `{data.tree}`.
 
 This is a general purpose format that fits generic hierarchical tree encoding needs. Each node of the tree is associated with predictor values through the `attributes` in the data `Node` object.
 
- - A very basic example is the `acme` dataset to show you how the two predictors values `cost` and `p` are associates attributes of each node in the hierarchy : 
+|{tabnet} concept| {data.tree} concept |see command|
+|---|---|---|
+|dataset predictor| Node `attributesAll` |  acme example |
+|dataset multi-label target| Node hierarchy | print(acme) |
+
+
+A very basic example is the `acme` dataset to show you how the two predictors values `cost` and `p` are associates attributes of each node in the hierarchy : 
 ```{r}
 data(acme, package = "data.tree")
 acme$attributesAll
 print(acme, "cost", "p" , limit = 8)
 ```
+So printing Node objects reverse the usual ordering, as target is printed first in column `levelName`, and predictors printed right of it. 
+
+As you can see, only leaf nodes of the tree gets predictors value. {tabnet} will take this into account via an `ancestor` square sparse tensor registering all possible parent-child relation among the target labels.
+
 
-- Multiple manual or programmatic methods are available to create or update predictors. They are detailled in the `vignette("data.tree", package = "data.tree")`.
+## Data preparation
+
+Multiple manual or programmatic methods are available to create or update predictors. They are detailled in the `vignette("data.tree", package = "data.tree")`.
 
 - a lot of native hierarchical data-format conversion from files to `Node` are covered by the`{data.tree}` package. You can find them in the "Create tree from a file" section of the same vignette. If needed, the `{ape}` package covers a lot of conversion format to the `philo` format. Thus you can reach the `Node` format in maybe two transformation steps...
 
@@ -74,25 +86,33 @@ As `as.Node()` will only consider the as.numeric() values of a factor(), you sho
 
 Your dataset hierarchy will be turn internally into multi-outcomes named `level_1` to `level_n`, n beeing the depth of your tree. Thus column names starting with `level_` should be avoided.
 
-### Ensure the last hierarchy of the tree is the observation id
+### Ensure the last hierarchy of the tree is the **observation id**
 
 The tree only keeps a single row of attributes per tree leaf. Thus in order to transfer your complete predictors dataset into the Node object, you must keep the last level of the hierarchy to be a unique observation identifier (last resort beeing `rowid_to_column()` to achieve it).
 
 The classification will be done **removing the last level of hierarchy** in any case.
 
-### Ensure there is a root level in the hierarchy
+### Ensure there is a **root level** in the hierarchy
 
 The tree should have a single root for all nodes to be consistent. Thus you have to use a constant prefix to all `pathString`.
 
 The classification will be done **removing the first level of hierarchy** in any case.
 
+### Ensure there is no **missing values** in the hierarchical classes
+
+Missing values should be replaced. Turning them  by an explicit "Unknown_something" is a good approach.
+
+
 Now let's have all those rules applied to the `starwars_tree` : 
 
 ```{r}
 # demonstration of reserved column modification in Node construction
 starwars_tree <- starwars %>% 
   rename(`_name` = "name", `_height` = "height") %>% 
-  mutate(pathString = paste("StarWars_characters", species, sex, `_name`, sep = "/")) %>%
+  mutate(
+    species = coalesce(species, "Unknown_Species"),
+    sex     = coalesce(sex, "Unknown_Sex"),
+    pathString = paste("StarWars_characters", species, sex, `_name`, sep = "/")) %>%
   as.Node()
 print(starwars_tree, "name", "_name","_height", "mass", "eye_color", limit = 8)
 ```