You can iterate through the model parameters and modify them as u want.
for n,p in model.named_parameters():
    if n == desired_layer_name:
       p.grad=your_func(*args)

if you want to modify all the gradients based in the loss’ gradient I think (cannot promise) you can do it with backward hooks. It basically runs some code once backward is called.
https://pytorch.org/tutorials/beginner/former_torchies/nnft_tutorial.html#forward-and-backward-function-hooks 1