clean up bert next_data (#13650)

train iter was designed to never stop for both real and fake data
This commit is contained in:
chenyu
2025-12-11 22:56:28 -05:00
committed by GitHub
parent ab2220b834
commit 01e9ad0d52

View File

@@ -1142,12 +1142,7 @@ def train_bert():
train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"])
pt = time.perf_counter()
try:
next_data = next(train_it)
except StopIteration:
next_data = None
next_data = next(train_it)
dt = time.perf_counter()
device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"