fixt print model only few layers

[jasnei]

model = max_vit_base_224(num_classes=10) print(model)

---

terminal output: (stem): Sequential( (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)) (1): GELU() (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (3): GELU() ) (head): Linear(in_features=768, out_features=10, bias=True) )

---

model = max_vit_base_224(num_classes=10) data = torch.rand(1, 3, 224, 224) from torchsummary import summary summary(model, data, device='cpu')

---

Layer (type:depth-idx) Output Shape Param #

---

├─Sequential: 1-1 [-1, 64, 112, 112] -- | └─Conv2d: 2-1 [-1, 64, 112, 112] 1,792 | └─GELU: 2-2 [-1, 64, 112, 112] -- | └─Conv2d: 2-3 [-1, 64, 112, 112] 36,928 | └─GELU: 2-4 [-1, 64, 112, 112] -- ├─Linear: 1-2 [-1, 10] 7,690

---

Total params: 46,410 Trainable params: 46,410 Non-trainable params: 0 Total mult-adds (M): 484.14

---

Input size (MB): 0.57 Forward/backward pass size (MB): 12.25 Params size (MB): 0.18 Estimated Total Size (MB): 13.00

---