(self, idx_in_batch, epoch_idx)
| 195 | self.accepts_arg = arg_count > 0 |
| 196 | |
| 197 | def _callback_args(self, idx_in_batch, epoch_idx): |
| 198 | if not self.accepts_arg: |
| 199 | return () |
| 200 | if idx_in_batch is not None: |
| 201 | arg = _types.SampleInfo( |
| 202 | self._current_sample + idx_in_batch, idx_in_batch, self._current_iter, epoch_idx |
| 203 | ) |
| 204 | elif self._batch_info: |
| 205 | arg = _types.BatchInfo(self._current_iter, epoch_idx) |
| 206 | else: |
| 207 | arg = self._current_iter |
| 208 | return (arg,) |
| 209 | |
| 210 | def _get_batch(self, epoch_idx): |
| 211 | try: |