rensjochemsen

I encountered the same error and tried to fix it your way. This is what I had: ``` opt_maker = lambda x: SWA(optim.AdamW(x.parameters(), lr=self.learning_rate, eps=1e-4)) if self.swa \ else optim.AdamW(x.parameters(),...