Resetting internal state of optimizer

C++
1

I would like to reset the internal state of the optimizer after a couple of epochs. I figured out that I can do

opt.state().clear();

but afterwards the call

opt.step(closure);

with closure being defined as

auto closure = [&]() {
        // Reset gradients
        net_->zero_grad();

        // Execute the model on the inputs
        outputs = net_->forward(inputs);

        // Compute the loss value
        loss = this->loss(outputs, epoch);

        // Compute gradients of the loss w.r.t. the model parameters
        loss.backward({}, true, false);

        return loss;
      };

gives an immediate segmentation fault.

Any help is appreciated.