PiBddlZddlmZdejdejddfdZdS)N)nnmodelscalerreturncd}|D]7}|s|j}||}|j|xj|zc_8dS)a Utility to scale the gradients of a model. This is useful for gradient accumulation where we want to normalize the gradients by the total number of tokens seen. Inputs: model (nn.Module): model whose gradients should be scaled scaler (torch.Tensor): scaling factor to apply to the gradients Outputs: None (grad fields are modified in place) N) parametersdevicetograd)rrr ps s/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/torchtune/training/_grad_scaler.py scale_gradsr sgF      'XFYYv&&F 6  FFf FF )torchrModuleTensorrrr rsY ry%,4r