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