def get_batch(data, batch_size, seq_len, index):
  """Batch the data."""
  elems_in_batch = batch_size * seq_len
  batches_count = len(data) // elems_in_batch

  batch_start = elems_in_batch * (index % batches_count)
  batch_end = batch_start + elems_in_batch

  batch = data[batch_start:batch_end]
  batch = batch.reshape(batch_size, seq_len)

  return batch