From e298ac4c453ee9798a41d8a357e54348473b36bd Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Wed, 5 Feb 2025 21:27:42 +0100
Subject: [PATCH] manually loading the correct modle again

---
 src/aiunn/finetune.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py
index 4fde964..7956384 100644
--- a/src/aiunn/finetune.py
+++ b/src/aiunn/finetune.py
@@ -138,7 +138,7 @@ def main():
     val_parquet_path = "/root/training_data/vision-dataset/image_vec_upscaler.parquet"
                 
     # Load pretrained model
-    model = AIIA.load("/root/vision/AIIA/AIIA-base-512")
+    model = AIIABase.load("/root/vision/AIIA/AIIA-base-512")
     
     # Add final upsampling layer if needed (depending on your specific architecture)
     if hasattr(model, 'chunked_'):