ViT(
(to_patch_embedding): Sequential(
(0): Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=32, p2=32)
(1): LayerNorm((3072,), eps=1e-05, elementwise_affine=True)
(2): Linear(in_features=3072, out_features=1024, bias=True)
(3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
)
(dropout): Dropout(p=0.1, inplace=False)
(transformer): Transformer(
(norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
(layers): ModuleList(
(0-5): 6 x ModuleList(
(0): Attention(
(norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
(attend): Softmax(dim=-1)
(dropout): Dropout(p=0.1, inplace=False)
(to_qkv): Linear(in_features=1024, out_features=3072, bias=False)
(to_out): Sequential(
(0): Linear(in_features=1024, out_features=1024, bias=True)
(1): Dropout(p=0.1, inplace=False)
)
)
(1): FeedForward(
(net): Sequential(
(0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
(1): Linear(in_features=1024, out_features=2048, bias=True)
(2): GELU(approximate='none')
(3): Dropout(p=0.1, inplace=False)
(4): Linear(in_features=2048, out_features=1024, bias=True)
(5): Dropout(p=0.1, inplace=False)
)
)
)
)
)
(to_latent): Identity()
(mlp_head): Linear(in_features=1024, out_features=1000, bias=True)
)
