@osalpekar Thanks for your reply.
The above code indeed works when an error occurs during the forward pass, skipping a batch and happily continuing, but when the error occurs in the backward pass I get the exception above (tested by increasing batch size until a OOM occurs during in the forward pass).
In general, it works well, and I do not yet need to use model parallelization, I do get an OOM error (about once a day). To get a bit of an idea, this is for my MRI reconstruction project 1, where I can use 40GB of memory, the model is typically around ±34GB (batch size 1 per GPU), but I can get an OOM in the backward pass.
Not particularly sure why and how this happens, but seems to be rather deep in the pytorch internals, and checking how they solve it in e.g. detectron2 it seems like a pragmatic way to solve it this way.
Blatantly ignoring the above exception and continue with the next batch just freezes the training by the way, so there should be something I need to reset,