Struct Flatten 

pub struct Flatten;

Flattens the input tensor into a 2D tensor (batch_size, remaining_features).

Assumes the first dimension is the batch dimension and flattens all subsequent dimensions. Input shape: (B, D1, D2, …) Output shape: (B, D1 * D2 * …)

Implementations

impl Flatten

pub fn new() -> Self

Trait Implementations

impl Default for Flatten

fn default() -> Self

Returns the “default value” for a type.

impl Module for Flatten

fn forward(&self, x: &Tensor) -> Tensor

fn parameters(&self) -> Vec<Tensor>

fn state_dict(&self) -> StateDict

fn load_state_dict(&mut self, _state: &StateDict)

fn zero_grad(&mut self)

fn to_device(&mut self, device: Device)

Move all module parameters to a specific device

fn train(&mut self, _mode: bool)

Switch between training and evaluation modes. Important for layers like BatchNorm and Dropout.

fn eval(&mut self)