Hmm, I guess someone needs to submit a pull request to update fast.ai docs+library to make this easier but check out my previous posts: How to review images for bad classifications + How do I get my list of predictions match the order of the images in my test folder?
It’s not a copy&paste solution but I think these posts should provide you with all the ingredients. I guess the main point for you is that you the need the ds_idx instead of of just iterating of .items.