make autograd.grad at DDP model

Can I use torch.autograd.grad between loss and apex DDP model.parameters during training(after wrapping the model with DDP)?

I'm working on implement gradNorm on apex DDP model.
For that, I need to get a gradient between loss and model parameters each iteration.

When I use multi-GPU, my code is stuck at the time when it accesses model parameter.
I'm not sure it is caused by accessing model parameter or autograd.
At using single GPU, everything is working fine except training time. :)

If autograd is not working with apex DDP model (same as torch DDP),
Or if is there any way to using autograd between loss and model parameter during training time,
Please let me know, Thanks.

The issue names no source file or test. Start by reproducing the multi-GPU hang with torch.autograd.grad on an apex DDP model and model.parameters, then inspect the apex DDP entry point and compare it with single-GPU behavior. Done means identifying whether parameter access or autograd causes the hang and documenting a confirmed compatible approach or a reproducible failure.