Merge pull request #498 from JuliaTrustworthyAI/497-bug-in-train_test_split-function

fixed

Author: pat-alt

CHANGELOG.md

@@ -10,12 +10,13 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ### Changed 
 
-- Slight changes to the implementation of `ProbeGenerator` (no longer calling a redundant `hinge_loss` function for all other generators). 
+- Addressed bug in `train_test_split` function. [#497]
+- Slight changes to the implementation of `ProbeGenerator` (no longer calling a redundant `hinge_loss` function for all other generators). [#492]
 
 ### Added
 
-- Added a warning message to the `ProbeGenerator` pointing to the issues with with current implementation.
-- Added links to papers to all docstrings for generators. 
+- Added a warning message to the `ProbeGenerator` pointing to the issues with with current implementation. [#492]
+- Added links to papers to all docstrings for generators. [#492]
 
 ## Version [1.3.5] - 2024-10-28
 

src/data_preprocessing/utils.jl

@@ -45,7 +45,7 @@ function train_test_split(data::CounterfactualData; test_size=0.2, keep_class_ra
         class_samples = [
             sample(
                 findall(vec(y .== cls)),
-                Int(floor(test_size * cls_ratio * n_per_class));
+                Int(floor(test_size * cls_ratio * N));
                 replace=false,
             ) for (cls, cls_ratio) in zip(classes_, class_ratios)
         ]