fixed saving

pyproject.toml

@@ -10,7 +10,7 @@ include = '\.pyi?$'
 
 [project]
 name = "aiia"
-version = "0.3.2"
+version = "0.3.3"
 description = "AIIA Deep Learning Model Implementation"
 readme = "README.md"
 authors = [

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = aiia
-version = 0.3.2
+version = 0.3.3
 author = Falko Habel
 author_email = falko.habel@gmx.de
 description = AIIA deep learning model implementation

src/aiia/__init__.py

@@ -4,4 +4,4 @@ from .data.DataLoader import DataLoader
 from .pretrain.pretrainer import Pretrainer, ProjectionHead
 
 
-__version__ = "0.3.2"
+__version__ = "0.3.3"

src/aiia/pretrain/pretrainer.py

@@ -246,7 +246,7 @@ class Pretrainer:
 
             if val_loss < best_val_loss:
                 best_val_loss = val_loss
-                self.model.save(output_path)
+                self.model.save_pretrained(output_path)
                 print("Best model saved!")
 
         losses_path = os.path.join(os.path.dirname(output_path), 'losses.csv')