Glad to know it wasn’t just me being dumb.
For the benefit of anyone who reads this later (https://xkcd.com/979/ 48) my hacky solution was to checkpoint all layers in the model. Not ideal, but it seems to get the job done, and the run time is not -too- bad