Automated Code Change

PiperOrigin-RevId: 724300701
google · Feb 7, 2025 · 22bfbc2 · 22bfbc2
1 parent 38fa93a
commit 22bfbc2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trax/data/tf_inputs.py b/trax/data/tf_inputs.py
@@ -186,7 +186,7 @@ def append_targets(example):
     # Skip a random fraction at the beginning of the stream.  The skip is
     # essential for synchronous highly-parallel training to avoid multiple
     # replicas reading the same data in lock-step.
-    dataset = dataset.skip(random.randint(0, _MAX_SKIP_EXAMPLES))
+    dataset = dataset.skip(random.randint(0, int(_MAX_SKIP_EXAMPLES)))
   dataset = preprocess_fn(dataset, training)
   dataset = dataset.shuffle(shuffle_buffer_size)
   return dataset.prefetch(8)