Skip to content

Commit

Permalink
add max_size to DataDomain and changed max_size default value to None…
Browse files Browse the repository at this point in the history
… instead -1 in ShimmerDataset
  • Loading branch information
leopoldmt committed Oct 15, 2024
1 parent 519d5f2 commit 6556d02
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 2 deletions.
4 changes: 2 additions & 2 deletions shimmer/data/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ def __init__(
dataset_path: str | Path,
split: str,
domain_classes: Mapping[DomainDesc, type[DataDomain]],
max_size: int = -1,
max_size: int | None = None,
transforms: Mapping[str, Callable[[Any], Any]] | None = None,
domain_args: Mapping[str, Any] | None = None,
):
Expand Down Expand Up @@ -94,7 +94,7 @@ def __init__(
lengths = {len(domain) for domain in self.domains.values()}
assert len(lengths) == 1, "Domains have different lengths"
self.dataset_size = next(iter(lengths))
if self.max_size != -1:
if self.max_size is not None:
assert (
self.max_size <= self.dataset_size
), "Max sizes can only be lower than actual size."
Expand Down
1 change: 1 addition & 0 deletions shimmer/data/domain.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ def __init__(
self,
dataset_path: str | Path,
split: str,
max_size: int | None = None,
transform: Callable[[Any], _T] | None = None,
additional_args: dict[str, Any] | None = None,
) -> None:
Expand Down

0 comments on commit 6556d02

Please sign in to comment.