Struct GraphClassificationHead 

pub struct GraphClassificationHead {
    pub pooling_layer: Tensor,
    pub classifier: Tensor,
    pub bias: Tensor,
}

Graph classification head

Fields

pooling_layer: Tensor

classifier: Tensor

bias: Tensor

Implementations

impl GraphClassificationHead

pub fn new( input_dim: usize, num_classes: usize, ) -> Result<Self, FoundationModelError>

Trait Implementations

impl Debug for GraphClassificationHead

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl TaskHead for GraphClassificationHead

fn forward(&self, embeddings: &Tensor) -> Result<Tensor, FoundationModelError>

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