class_docstr: <a id="ClassificationHead"></a>
  
  ## Classification Head
  
  They do average pooling (taking the mean of all patch embeddings) and a final linear transformation
  to predict the log-probabilities of the image classes.
class_docstr_tok: ["<","a","id=","''","ClassificationHead","''",">","<","\/a",">","#","#","Classification","Head","They","do","average","pooling","(","taking","the","mean","of","all","patch","embeddings",")","and","a","final","linear","transformation","to","predict","the","log-probabilities","of","the","image","classes","."]
docstr: * `d_model` is the number of channels in patch embeddings, $h$
  * `n_classes` is the number of classes in the classification task
docstr_tok: ["*","`","d_model","`","is","the","number","of","channels","in","patch","embeddings",",","$","h","$","*","`","n_classes","`","is","the","number","of","classes","in","the","classification","task"]
code: def __init__(self, d_model: int, n_classes: int):
      """
      * `d_model` is the number of channels in patch embeddings, $h$
      * `n_classes` is the number of classes in the classification task
      """
      super().__init__()
      # Average Pool
      self.pool = nn.AdaptiveAvgPool2d((1, 1))
      # Linear layer
      self.linear = nn.Linear(d_model, n_classes)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/conv_mixer/__init__.py
imports: import torch
  from torch import nn
  from labml_helpers.module import Module
  from labml_nn.utils import clone_module_list
label_desc: Use image node_id 1 to create a new ClassificationHead object from inherited base classes: Module with example: obj = ClassificationHead(d_model, n_classes)