How to use luz with custom optimizer?

[icejean]

Hi all,
I'm new to R torch and luz, according to issue#36, luz requires modules that are not yet initialized, but definding a custom optimizer require parameters of a model, so the model should be initialized, just as below:

model <- net()$to(device = device)        
optimizer <- optim_adamw(model$parameters, lr=0.001)
criterion <- nn_cross_entropy_loss()    

I can run the example of luz without custom optimizer, just as the following lines:

net <- nn_module(
......
)
fitted <- net %>%
  setup(
    loss = nn_cross_entropy_loss(),
    # optim_adam, optim_adamw, optim_rprop, optim_adadelta, optim_asgd, optim_rmsprop, 
    optimizer = optim_adam, 
    metrics = list(
      luz_metric_accuracy()
    )
  ) %>%
  fit(train_dl, epochs = 15, valid_data = test_dl)

So my question is, how can I use a custom optimizer together with luz?
for example:

model <- net()$to(device = device) 
# optim_adam, optim_adamw, optim_rprop, optim_adadelta, optim_asgd, optim_rmsprop
optimizer <- optim_adam(model$parameters, lr=0.001) 
criterion <- nn_cross_entropy_loss() 

fitted <- model %>%
  setup(
    loss = criterion,
    optimizer = optimizer, 
    metrics = list(
      luz_metric_accuracy()
    )
  ) %>%
  fit(train_dl, epochs = 15, valid_data = test_dl)

I'll get an error when I code in this way.

Error in x$get_inherit() : attempt to apply non-function

Any idea?
Thanks in advance.