Trait border_tch_agent::model::SubModel2

pub trait SubModel2 {
    type Config;
    type Input1;
    type Input2;
    type Output;

    // Required methods
    fn build(var_store: &VarStore, config: Self::Config) -> Self;
    fn clone_with_var_store(&self, var_store: &VarStore) -> Self;
    fn forward(
        &self,
        input1: &Self::Input1,
        input2: &Self::Input2
    ) -> Self::Output;
}

Neural network model that can be initialized with VarStore and configuration.

The difference from SubModel is that this trait takes two inputs.

Required Associated Types

type Config

Configuration from which SubModel2 is constructed.

type Input1

Input of the SubModel2.

type Input2

Input of the SubModel2.

type Output

Output of the SubModel2.

Required Methods

fn build(var_store: &VarStore, config: Self::Config) -> Self

Builds SubModel2 with VarStore and SubModel2::Config.

fn clone_with_var_store(&self, var_store: &VarStore) -> Self

Clones SubModel2 with VarStore.

fn forward(&self, input1: &Self::Input1, input2: &Self::Input2) -> Self::Output

A generalized forward function.

Implementors

impl SubModel2 for Mlp

type Config = MlpConfig

type Input1 = Tensor

type Input2 = Tensor

type Output = Tensor