From dd7d569834e20a48ea4c2e2b3172a866491a4ff4 Mon Sep 17 00:00:00 2001 From: Davis Wertheimer Date: Mon, 10 Feb 2025 15:34:45 -0500 Subject: [PATCH] diag print4 --- torchdata/stateful_dataloader/ibm_rescalable.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torchdata/stateful_dataloader/ibm_rescalable.py b/torchdata/stateful_dataloader/ibm_rescalable.py index 594f9f9c6..e2bb99168 100644 --- a/torchdata/stateful_dataloader/ibm_rescalable.py +++ b/torchdata/stateful_dataloader/ibm_rescalable.py @@ -493,6 +493,7 @@ def load_state_dict(self, state_dict): self.filesizes = file_info self.shard_states = shard_states[self.rank] else: + print("GOTHERE 3") shard_states = [s[0] for s in shard_states.split(1)] # [w] n 5 shard_states = torch.cat(shard_states, dim=0) # wn 5 # Sort shards by epoch count