burn_core/nn/cache/
base.rs

1use crate::tensor::Tensor;
2use crate::tensor::backend::Backend;
3
4pub(crate) enum CacheState<T> {
5    Value(T),
6    Empty,
7}
8
9/// A cache for a tensor.
10pub struct TensorCache<B: Backend, const D: usize> {
11    pub(crate) state: CacheState<Tensor<B, D>>,
12}
13
14impl<B: Backend, const D: usize> TensorCache<B, D> {
15    /// Creates a new empty cache.
16    ///
17    /// # Returns
18    ///
19    /// The empty cache.
20    pub fn empty() -> Self {
21        Self {
22            state: CacheState::Empty,
23        }
24    }
25}