Thank you for your reply!
As you pointed, I use davda54’s SAM implementation without closure.
And it works well.
But I use pytorch-lightning. In pytorch-lightning environment, I have to write some extra codes to calculate gradients two times. My motivation is to solve this problem. I want to use SAM as other optimizers.
I have understood that closure support is tricky. I decided to use SAM without closure. Thank!!