I have a module MyModel that I want to instantiate then call a custom method on it, plus some base methods, then return a generic ModuleHolder<Module> to it to the calling code.
What is the correct way to return a ModuleHolder<Module>?
ModuleHolder<Module> load_model(std::string w, std::string device) {
auto state_dict = load_weights(w);
auto module = MyModel(insize, outsize);
// call a custom method
module->load_state_dict(state_dict);
// call base methods
if (device == "gpu") {
module->to(device);
module->to(c10::kHalf);
}
return module; // How to return a ModuleHolder<Module> not ModuleHolder<MyModel>?
}
Example model definition:
struct ConvolutionImpl : Module {
ConvolutionImpl(int size = 1, int outsize = 1, int k = 1, int stride = 1) {
conv = register_module("conv", Conv1d(Conv1dOptions(size, outsize, k).stride(stride).padding(k/2)));
activation = register_module("activation", SiLU());
}
torch::Tensor forward(torch::Tensor x) {
return activation(conv(x));
}
Conv1d conv{nullptr};
SiLU activation{nullptr};
};
TORCH_MODULE(Convolution);
struct MyModelImpl : Module {
Model(int size, int outsize) {
conv1 = register_module("conv1", Convolution(1, 4, 5, 1));
conv2 = register_module("conv2", Convolution(1, 16, 5, 1));
conv3 = register_module("conv3", Convolution(16, outsize, 19, 5));
}
void load_state_dict(std::vector<torch::Tensor> weights) {
assert (weights.size() == parameters().size());
for (size_t idx = 0; idx < weights.size(); idx++) {
parameters()[idx].data() = weights[idx].data();
}
}
Convolution conv1{nullptr}, conv2{nullptr}, conv3{nullptr};
};
TORCH_MODULE(MyModel);