Modify torch::jit::script::Module Parameters

Hey y’all!

I am using LibTorch v1.12.1 and I have been, unsuccessfully, trying to modify the parameters of a TorchScript module that I have exported from Python and then loaded into C++.

First what I have in one section of the app is:

    auto m_module = torch::jit::load(modelPath);

    auto parameters = std::vector<torch::Tensor>{};

    for (const auto &param : m_module.parameters()) {
       parameters.push_bach(param.detach().clone());
    }

The app then modifies the parameters in the std::vector in some way say:

for (auto p : parameters) {
    p += 10;
}

Then in another section of the code I have:

    auto idx= std::size_t{0};
    for (auto lhsParam : m_module.named_parameters()) {
        m_module.register_parameter(lhsParam.name, parameters.at(idx), true);
        ++i;
    }

It seems like the register_parameter function isn’t actually updating the parameters of the m_module.

Is there something I’m missing here? Or does the register_parameter function not work that way? Is there another way of updating the parameters of m_module?

Any tips would be appreciated.

Thanks!

Found a good solution:

     // ! Disables autograd which makes parameter copying possible
    torch::autograd::GradMode::set_enabled(false);

    auto idx= std::size_t{0};
    for (auto lhsParam : m_module.named_parameters()) {
        lhsParam.value.copy_(parameters.at(idx));
        ++i;
    }
    
    torch::autograd::GradMode::set_enabled(true);

Hi Esteban, did you also managed to extract weights from such jit module and import them into a torch module?