How can I save weights/model in stratified K fold cross validation?

I am dealing with IDRID disease grading which is a small & imbalanced and also my architecture is complex, so I am using k-fold here. But I am not sure how and at what point should I save weights. Please help!

What is your final objective? To ensemble?

also my architecture is complex, so I am using k-fold here

Why so? You generally use k-fold to gauge robustness/generizability, not much to do with the complexity of the architecture right?

so i’ll have 5 files for 5 folds this way, how to load these weights during eval in code?