VitExitLayersΒΆ

Exit block of ViT.

Abstract Signature:

VitExitLayers(hidden_dim: int, output_dim: int)

PyTorch

API: β€”
Strategy: Custom / Partial

PaxML / Praxis

API: praxis.layers.VitExitLayers
Strategy: Direct Mapping