Can't pickle SWA: AttributeError: 'SWA' object has no attribute 'defaults' about contrib

[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(), lr=self.learning_rate, eps=1e-4)
self.optimizer = opt_maker(self.network)

After fix:

opt_maker = lambda x: SWA(optim.AdamW(x.parameters(), lr=self.learning_rate, eps=1e-4)) if self.swa \
                                else optim.AdamW(x.parameters(), lr=self.learning_rate, eps=1e-4)
self.optimizer = opt_maker(self.network)
self.optimizer.param_groups = self.optimizer.optimizer.param_groups
self.optimizer.state = self.optimizer.optimizer.state
self.optimizer.defaults=self.optimizer.optimizer.defaults

Maybe this was not what you meant. For me it resulted in the the error:

BrokenPipeError: [Errno 32] Broken pipe