example.py

@@ -7,7 +7,7 @@ config = AIIAConfig(model_name="AIIA-Base-512x20k")
 model = AIIABase(config)
 
 # Initialize pretrainer with the model
-pretrainer = Pretrainer(model, learning_rate=1e-4, config=config)
+pretrainer = Pretrainer(model, learning_rate=1e-4)
 
 # List of dataset paths
 dataset_paths = [

requirements.txt

@@ -2,7 +2,4 @@ torch>=2.5.0
 numpy
 tqdm
 pytest
-pillow
-pandas
-torchvision
-pyarrow
+pillow

src/aiia/model/Model.py

@@ -130,6 +130,7 @@ class AIIABaseShared(AIIA):
         self.max_pool = nn.MaxPool2d(
             kernel_size=1,
             stride=1,
+            padding=1
         )
 
     def forward(self, x):