better error message for nn_head (#369)

sebffischer · web-flow · commit addf4eeff935 · 2025-03-25T10:53:10.000+01:00
diff --git a/R/PipeOpTorchHead.R b/R/PipeOpTorchHead.R
@@ -36,7 +36,9 @@ PipeOpTorchHead = R6Class("PipeOpTorchHead",
   ),
   private = list(
     .shapes_out = function(shapes_in, param_vals, task) {
-      assert_true(length(shapes_in[[1]]) == 2L)
+      if (length(shapes_in[[1]]) != 2L) {
+        stopf("PipeOpTorchHead expects 2D input, but got %s.", shape_to_str(shapes_in))
+      }
       d = get_nout(task)
       list(c(shapes_in[[1]][[1]], d))
     },
diff --git a/tests/testthat/test_PipeOpTorchHead.R b/tests/testthat/test_PipeOpTorchHead.R
@@ -12,3 +12,10 @@ test_that("PipeOpTorchHead paramtest", {
   res = expect_paramset(po_test, torch::nn_linear, exclude = c("out_features", "in_features"))
   expect_paramtest(res)
 })
+
+
+test_that("correct error message", {
+  task = nano_imagenet()
+  graph = po("torch_ingress_ltnsr") %>>% po("nn_head")
+  expect_error(graph$train(task), "expects 2D input")
+})
