Set's up a
nn_module to use with luz
The setup function is used to set important attributes and method for
to be used with luz.
nn_modulethat you want set up.
function, optional) An optional function with the signature
function(input, target). It's only requires if your
nn_moduledoesn't implement a method called
torch_optimizer, optional) A function with the signature
function(parameters, ...)that is used to initialize an optimizer given the model parameters.
list, optional) A list of metrics to be tracked during the training procedure.
function) A functions that takes the loss scalar values as it's parameter. It must call
torch::autograd_backward(). In general you don't need to set this parameter unless you need to customize how luz calls the
backward(), for example, if you need to add additional arguments to the backward call. Note that this becomes a method of the
nn_modulethus can be used by your custom
step()if you override it.
A luz module that can be trained with