From 8c9853ade32055f06a1d6c0ce224740179596f1a Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Thu, 24 Apr 2025 09:37:19 +0200
Subject: [PATCH] fixed saving

---
 pyproject.toml                  | 2 +-
 setup.cfg                       | 2 +-
 src/aiia/__init__.py            | 2 +-
 src/aiia/pretrain/pretrainer.py | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8975d57..4b8bd0e 100644
--- a/pyproject.toml
+++ b/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 = [
diff --git a/setup.cfg b/setup.cfg
index 65869f3..202e051 100644
--- a/setup.cfg
+++ b/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
diff --git a/src/aiia/__init__.py b/src/aiia/__init__.py
index 3a8a200..9a9e138 100644
--- a/src/aiia/__init__.py
+++ b/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"
diff --git a/src/aiia/pretrain/pretrainer.py b/src/aiia/pretrain/pretrainer.py
index 62c01c0..9cde4b8 100644
--- a/src/aiia/pretrain/pretrainer.py
+++ b/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')