Saving train state of step 10000

run_distillation.py

@@ -355,7 +355,7 @@ class DataTrainingArguments:
         },
     )
     prefetch_size: int = field(
-        default=None,
+        default=0,
         metadata={"help": "Number of samples to pre-fetch if using an iterable dataset."},
     )
     timestamp_probability: float = field(
@@ -603,13 +603,16 @@ def get_data_loader(
     if skip_batches > 0:
         dataset = dataset.skip(skip_batches * batch_size)
 
+    if prefetch_size > 0:
+        dataset = IterableWrapper(dataset)
+        dataset = dataset.prefetch(prefetch_size)
+
     data_loader = DataLoader(
         dataset,
         batch_size=batch_size,
         drop_last=drop_last,
         pin_memory=pin_memory,
         collate_fn=data_collator,
-        prefetch_factor=prefetch_size,
         num_workers=dataloader_num_workers,
     )
 

runs/events.out.tfevents.1704385834.t1v-n-d928564b-w-0.141489.0.v2

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:71d0855527f0576381a5d8ed96a21fe622bb820a43e128303df870acc54648ad
-size 78

runs/{events.out.tfevents.1704385679.t1v-n-d928564b-w-0.137863.0.v2 → events.out.tfevents.1704446248.t1v-n-d928564b-w-0.202887.0.v2}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f2edf41629445d3e3849683885e19a83ce2f22ae2e8a0474d0a5dd35e4f45fcc
+oid sha256:52802eacb6c8b7f0ad8042620d33aa0abb6b7910a906fb6545a3189673cdb17e
 size 78