Struct NodeClassificationHead 

pub struct NodeClassificationHead {
    pub classifier: Tensor,
    pub bias: Tensor,
}

Node classification head

Fields

classifier: Tensor

bias: Tensor

Implementations

impl NodeClassificationHead

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

Trait Implementations

impl Debug for NodeClassificationHead

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

Formats the value using the given formatter.

impl TaskHead for NodeClassificationHead

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

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