Method __init__

axlearn/common/layers.py:1187–1190

(self, cfg: Config, *, parent: Module)

        super().__init__(cfg, parent=parent)
        cfg = self.config
        self._add_child("metric", cfg.metric.set(num_classes=cfg.num_classes))

    def forward(self, input_batch: NestedTensor) -> Tensor:
        """Produces predictions for the given inputs.

        Args:
            input_batch: A dict with the following entries: