Skip to main content

burn_train/
components.rs

1use crate::{InferenceStep, TrainStep};
2use burn_core::module::Module;
3
4/// A single keyword trait for a Burn [module](burn_core::module::Module) used for learning.
5pub trait LearnerModel: TrainStep + InferenceStep + Module + core::fmt::Display + 'static {}
6
7impl<T> LearnerModel for T where T: TrainStep + InferenceStep + Module + core::fmt::Display + 'static
8{}
9
10/// Type for training input.
11pub(crate) type TrainingModelInput<M> = <M as TrainStep>::Input;
12/// Type for inference input.
13pub(crate) type InferenceModelInput<M> = <M as InferenceStep>::Input;
14/// Type for training output.
15pub(crate) type TrainingModelOutput<M> = <M as TrainStep>::Output;
16/// Type for inference output.
17pub(crate) type InferenceModelOutput<M> = <M as InferenceStep>::Output;