diff --git a/torchdata/datapipes/iter/util/bz2fileloader.py b/torchdata/datapipes/iter/util/bz2fileloader.py index 442df0392..7a45a8278 100644 --- a/torchdata/datapipes/iter/util/bz2fileloader.py +++ b/torchdata/datapipes/iter/util/bz2fileloader.py @@ -53,7 +53,7 @@ def __iter__(self) -> Iterator[Tuple[str, BufferedIOBase]]: pathname, data_stream = data try: extracted_fobj = bz2.open(data_stream, mode="rb") # type: ignore[call-overload] - new_pathname = pathname.rstrip(".bz2") + new_pathname = pathname.removesuffix(".bz2") yield new_pathname, StreamWrapper(extracted_fobj, data_stream, name=new_pathname) # type: ignore[misc] except Exception as e: warnings.warn(f"Unable to extract files from corrupted bzip2 stream {pathname} due to: {e}, abort!")