Model trained by AWC style cannot be saved

The model trained by neighbor_allreduce optimizer and awc communication style cannot be saved with the command "torch.save()". See the error information below.

The possible reason might be in make_hook function in optimizer.py

==================
File "XXX.py", line 227, in
torch.save(model, tmp_path)
File "/usr/local/lib/python3.7/dist-packages/torch/serialization.py", line 328, in save
_legacy_save(obj, opened_file, pickle_module, pickle_protocol)
File "/usr/local/lib/python3.7/dist-packages/torch/serialization.py", line 401, in _legacy_save
pickler.dump(obj)
AttributeError: Can't pickle local object '_DistributedReduceOptimizer._make_hook..hook'

Start in optimizer.py at the _DistributedReduceOptimizer.make_hook function and reproduce the failure with the neighbor_allreduce optimizer, AWC communication style, and torch.save(model, tmp_path). Trace the generated hook involved in serialization; done means the trained model can be saved without the local-object pickle error.

Written by the indexing model from the issue text.