model in rust, optimizer.step in python

[xnerhu]

Hello.

Is it possible, to create a model in rust using tch-rs, export it using pyo3, so it is accessible from python side (forward function), then call it from python with some tensors and do optimization part on the python side too? I see that it may be impossible due to the usage of nn::Path / VarStore.

pub type Model = Box<dyn Fn(&Tensor) -> (Tensor)>;

#[pyfunction(name = "create_rs_model")]
pub fn model(p: &nn::Path) -> Model {
    let layers = nn::seq()
        .add(nn::linear(p / "layer1", 2, 16, Default::default()))
        .add_fn(|xs| xs.relu())
        .add(nn::linear(p / "layer2", 16, 1, Default::default()));
    return Box::new(move |xs: &Tensor| layers.forward(xs));
}

def class Encoder(torch.nn.Module)
...

encoder  = Encoder()
rs_model = create_rs_model()

encoded = encoder(inputs)
out = rs_model(encoded)

loss = F...(out, true)
loss.backward()
optimizer.step()