I think you may have an off-by-one error here. if `len(keys) == len(batch)` you're going to get an IndexError here

why not just do keys[-1]