example.py

@@ -11,7 +11,7 @@ from torchvision import transforms
 
 
 class UpscaleDataset(Dataset):
-    def __init__(self, parquet_files: list, transform=None, samples_per_file=2500):
+    def __init__(self, parquet_files: list, transform=None, samples_per_file=5000):
         combined_df = pd.DataFrame()
         for parquet_file in parquet_files:
             # Load a subset from each parquet file

src/aiunn/finetune/trainer.py

@@ -260,7 +260,8 @@ class aiuNNTrainer:
             torch.cuda.empty_cache()
             
             # Check early stopping
-            if early_stopping(val_loss if self.validation_loader else avg_train_loss):
+            early_stopping(val_loss if self.validation_loader else avg_train_loss)
+            if early_stopping.early_stop:
                 print(f"Early stopping triggered at epoch {epoch + 1}")
                 break