[This](https://github.com/callummcdougall/ARENA_2.0/blob/4cda66b64b48dbabdc6b0bd6f5d7a86eea375507/chapter1_transformers/exercises/part3_indirect_object_identification/ioi_dataset.py#L532) file has several nicer methods for easy IOI usage
