From 8bbf5c6dc3858771fcf5770faa56502a9f29f0c7 Mon Sep 17 00:00:00 2001
From: Moritz Gunz <moritz.gunz@gmail.com>
Date: Tue, 4 Feb 2025 14:10:44 +0100
Subject: [PATCH] DFDataset: do not pickle sharding info, but rely on
 `_distrib_info` instead

Fixes #1678
---
 returnn/datasets/distrib_files.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/returnn/datasets/distrib_files.py b/returnn/datasets/distrib_files.py
index f43a91c0e..be9e48ba7 100644
--- a/returnn/datasets/distrib_files.py
+++ b/returnn/datasets/distrib_files.py
@@ -130,6 +130,8 @@ def get_sub_epoch_dataset(files_subepoch: List[Tuple[str, str]]) -> Dict[str, An
     https://github.com/rwth-i6/returnn/issues/1524.
     """
 
+    _getnewargs_exclude_attrs = CachedDataset2._getnewargs_exclude_attrs.union(("_num_shards", "_shard_index"))
+
     def __init__(
         self,
         *,