library(mlr3torch)

cb = torch_callback("loss_fn",
  state_dict = function() {
    self$ctx$loss_fn
  },
  load_state_dict = function(state_dict) {
    self$x = x
  }
)

learner = lrn("classif.mlp", neurons = 100, batch_size = 32, epochs = 10,
  callbacks = cb)

tsk_binary = tsk("iris")$filter(1:100)$droplevels()
tsk_multi = tsk("iris")

learner$train(tsk_binary)
learner$network(torch_randn(1, 4))
#> torch_tensor
#> -0.3322
#> [ CPUFloatType{1,1} ][ grad_fn = <AddmmBackward0> ]
class(learner$model$callbacks$loss_fn)
#> [1] "nn_bce_with_logits_loss" "nn_loss"                
#> [3] "nn_module"

learner$train(tsk_multi)
learner$network(torch_randn(1, 4))
#> torch_tensor
#>  1.0422  0.3267 -0.5658
#> [ CPUFloatType{1,3} ][ grad_fn = <AddmmBackward0> ]
class(learner$model$callbacks$loss_fn)
#> [1] "nn_cross_entropy_loss" "nn_weighted_loss"      "nn_loss"              
#> [4] "nn_module"

glrn = po("torch_ingress_num") %>>%
  nn("linear", out_features = 2) %>>%
  po("torch_loss", loss = "cross_entropy") %>>%
  po("torch_optimizer", "adamw") %>>%
  po("torch_model_classif", epochs = 1, batch_size = 32) |> as_learner()

glrn$train(tsk("sonar"))
#> Error in (function (self, target, weight, pos_weight, reduction) : output with shape [32, 1] doesn't match the broadcast shape [32, 2]
#> Exception raised from mark_resize_outputs at /Users/runner/work/libtorch-mac-m1/libtorch-mac-m1/pytorch/aten/src/ATen/TensorIterator.cpp:1208 (most recent call first):
#> frame #0: c10::Error::Error(c10::SourceLocation, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char>>) + 52 (0x10422811c in libc10.dylib)
#> frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char>> const&) + 140 (0x104224d6c in libc10.dylib)
#> ... 
#> This happened PipeOp torch_model_classif's $train()