burn-train 0.22.0-pre.1

Training crate for the Burn framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

use crate::{InferenceStep, TrainStep};
use burn_core::module::AutodiffModule;

/// A single keyword trait for a Burn [module](burn_core::module::Module) used for learning.
pub trait LearnerModel:
    TrainStep + InferenceStep + AutodiffModule + core::fmt::Display + 'static
{
}

impl<T> LearnerModel for T where
    T: TrainStep + InferenceStep + AutodiffModule + core::fmt::Display + 'static
{
}

/// Type for training input.
pub(crate) type TrainingModelInput<M> = <M as TrainStep>::Input;
/// Type for inference input.
pub(crate) type InferenceModelInput<M> = <M as InferenceStep>::Input;
/// Type for training output.
pub(crate) type TrainingModelOutput<M> = <M as TrainStep>::Output;
/// Type for inference output.
pub(crate) type InferenceModelOutput<M> = <M as InferenceStep>::Output;