I am currently trying to using apex with SWA like so:
model, optimizer = amp.initialize(model, optimizer, opt_level='O2')
model = DDP(model)

swa_model = torch.optim.swa_utils.AveragedModel(model)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)
swa_start = 6
swa_scheduler = SWALR(optimizer, swa_lr=0.05 * world_size)

At this line (swa_model = torch.optim.swa_utils.AveragedModel(model)) I am getting the following error:
Traceback (most recent call last):
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 19, in _wrap
    fn(i, *args)
  File "/home/jupyter/Flood_Comp/starter.py", line 249, in train
    swa_model = torch.optim.swa_utils.AveragedModel(model)
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/optim/swa_utils.py", line 89, in __init__
    self.module = deepcopy(model)
  File "/opt/conda/lib/python3.7/copy.py", line 169, in deepcopy
    rv = reductor(4)
  File "/opt/conda/lib/python3.7/site-packages/apex/parallel/distributed.py", line 271, in __getstate__
    del attrs['self.bucket_streams']
KeyError: 'self.bucket_streams'

Any pointers on mitigating this would be helpful.