fix(jit): jit_trace now works with batch-norm (#355)

Otherwise, the variance is an nan tensor during tracing which
-- for some reason -- does not get updated.

Resolves Issue #354

NEWS.md

@@ -2,6 +2,7 @@
 
 * fix: `LearnerTorchModel` can now be parallelized and trained with
   encapsulation activated.
+* fix: `jit_trace` now works in combination with batch normalization.
 
 # mlr3torch 0.2.0
 

R/utils.R

@@ -238,10 +238,16 @@ CallbacksNone = function() {
 
 get_example_batch = function(dl) {
   ds = dl$dataset
+  if (length(ds) < 2) {
+    stopf("Dataset needs to contain at least 2 observations")
+  }
   if (!is.null(ds$.getbatch)) {
-    ds$.getbatch(1)
+    ds$.getbatch(1:2)
   } else {
-    ds$.getitem(1)$unsqueeze(1)
+    torch_stack(list(
+      ds$.getitem(1),
+      ds$.getitem(2)
+    ))
   }
 }
 

tests/testthat/test_PipeOpTorchBatchNorm.R

@@ -40,3 +40,16 @@ test_that("PipeOpTorchBatchNorm3D paramtest", {
   res = expect_paramset(po("nn_batch_norm3d"), nn_batch_norm3d, exclude = "num_features")
   expect_paramtest(res)
 })
+
+test_that("jit_trace works (#354)", {
+  graph = po("torch_ingress_num") %>>%
+    nn("batch_norm1d") %>>%
+    nn("head") %>>%
+    po("torch_loss", t_loss("cross_entropy")) %>>%
+    po("torch_optimizer", t_opt("adamw")) %>>%
+    po("torch_model_classif", epochs = 1, batch_size = 50)
+  lrn = as_learner(graph)
+  task = tsk("iris")
+  lrn$train(task)
+  expect_prediction(lrn$predict(task))
+})