Bugfix: n_batch should always be <= n_ctx