code: def _load_linear_head(
      num_features: int, num_classes: int
  ) -> nn.Module:
      """Loads a linear head.
  
      Args:
          num_features: Number of input features.
          num_classes: Number of output classes.
  
      Returns:
          nn.Module: Linear head.
  
      """
      return nn.Linear(num_features, num_classes)
  
imports: from functools import partial
  from torch import nn
  from flash.core.registry import FlashRegistry
label_desc: Use image node_id 1 for calling a global function with example usage: _load_linear_head(num_features, num_classes) and returns: nn