Loads the next batch with sharding hint
(self, *args, **kwargs)
| 69 | return self.last_batch |
| 70 | |
| 71 | def load_next_batch(self, *args, **kwargs): |
| 72 | """Loads the next batch with sharding hint""" |
| 73 | return maybe_shard_with_name( |
| 74 | self.load_next_batch_pre_sharding(), |
| 75 | self.input_data_shardings, |
| 76 | self.config.shard_mode, |
| 77 | ) |
| 78 | |
| 79 | def check_example_batch(self): |
| 80 | if self.config.max_checkify: |