Just thought I’d throw this down here to help you guys out, I did do something like this already in fastai v2 I’m not 100% sure if it still works today, but you should steal the concept around it as it’s fairly straightforward