Fix: ensure correct handling of val_set_size as float or int (#1655)

* Fix: ensure correct handling of val_set_size as float or int

* chore: lint

---------

Co-authored-by: Wing Lian <wing.lian@gmail.com>
This commit is contained in:
Davide Caroselli
2024-05-28 18:00:32 +02:00
committed by GitHub
parent f5febc729a
commit 6a5a725f10

View File

@@ -494,7 +494,9 @@ def load_prepare_datasets(
test_fingerprint = md5(to_hash_test)
dataset = dataset.train_test_split(
test_size=cfg.val_set_size,
test_size=int(cfg.val_set_size)
if cfg.val_set_size == int(cfg.val_set_size)
else cfg.val_set_size,
shuffle=False,
seed=cfg.seed or 42,
train_new_fingerprint=train_fingerprint,