IllegalWork error occurs when using the latest PyTorch
**Describe the bug**

Hi, I am trying to use a latest PyTorch (build from source from the master branch). But when I run the NeMo training script, I meet with the following error:
```
File "/usr/local/lib/python3.10/dist-packages/apex/contrib/optimizers/distributed_fused_adam.py", line 1602, in _local_grad_norm
    self.grad_sync()
  File "/usr/local/lib/python3.10/dist-packages/apex/contrib/optimizers/distributed_fused_adam.py", line 1572, in grad_sync
    self._force_bucket_grad_sync()
  File "/usr/local/lib/python3.10/dist-packages/apex/contrib/optimizers/distributed_fused_adam.py", line 1241, in _force_bucket_grad_sync
    self._finish_bucket_grad_sync()
  File "/usr/local/lib/python3.10/dist-packages/apex/contrib/optimizers/distributed_fused_adam.py", line 1394, in _finish_bucket_grad_sync
    bucket.sync_wait()
  File "/usr/local/lib/python3.10/dist-packages/apex/contrib/optimizers/distributed_fused_adam.py", line 361, in sync_wait
    self.sync_request.wait()
  File "/workspace/pytorch/torch/distributed/distributed_c10d.py", line 1650, in __getattribute__
    raise RuntimeError(f"Illegal to call {name} on IllegalWork object")
RuntimeError: Illegal to call wait on IllegalWork object
```
I have found that this error is related to `distributed_fused_adam` under the data parallel setting. When I change the optimizer to `fused_adam`, the training script can run successfully.

My training script can run successfully under the default PyTorch from nvcr.io/nvidia/pytorch:23.06-py3


**Environment details**

My docker is based on nvcr.io/nvidia/pytorch:23.06-py3
PyTorch: build from source on the master branch (commit: e471c12a016a46d7b8d1c8c83a38021f08fce078)
Apex commit: 4e1ae43f7f7ac69113ef426dd15f37123f0a2ed3

I am also trying to update apex to 7995de18677295c5edeeab082179edbfdb6ee16a following README but seemingly it is not correctly built and it would raise `ValueError: Cannot resolve optimizer 'distributed_fused_adam'.` I have reflected this issue on https://github.com/NVIDIA/apex/issues/1722