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

@@ -3,6 +3,3 @@ numpy
 tqdm
 pytest
 pillow
-pandas
-torchvision
-pyarrow

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):