docstr: Implements the NEFTune forward pass for the model using forward hooks. Note this works only for
  torch.nn. Embedding layers. This method is slightly adapted from the original source code that can be found
  here: https://github.com/neelsjain/NEFTune.
  
  The input tensor is ignored since the noise is added to the output of the embedding layer.
  
  Returns:
      torch.Tensor: The output tensor from the module.
docstr_tok: ["Implements","the","NEFTune","forward","pass","for","the","model","using","forward","hooks",".","Note","this","works","only","for","torch.nn",".","Embedding","layers",".","This","method","is","slightly","adapted","from","the","original","source","code","that","can","be","found","here",":","https",":","\/\/github.com\/neelsjain\/NEFTune",".","The","input","tensor","is","ignored","since","the","noise","is","added","to","the","output","of","the","embedding","layer",".","Returns",":","torch.Tensor",":","The","output","tensor","from","the","module","."]
code: def hook_fn(
      self,
      module: torch.nn.Module,
      input: torch.Tensor,
      output: torch.Tensor,
  ) -> torch.Tensor:
      """Implements the NEFTune forward pass for the model using forward hooks. Note this works only for
      torch.nn. Embedding layers. This method is slightly adapted from the original source code that can be found
      here: https://github.com/neelsjain/NEFTune.
  
      The input tensor is ignored since the noise is added to the output of the embedding layer.
  
      Returns:
          torch.Tensor: The output tensor from the module.
      """
      if module.training:
          dims = torch.tensor(output.size(1) * output.size(2))
          mag_norm = module.neftune_noise_alpha / torch.sqrt(dims)
          output = output + torch.zeros_like(output).uniform_(
              -mag_norm, mag_norm
          )
      return output
  
file_path: ludwig/ludwig/modules/training_hooks.py
imports: import logging
  from abc import ABC, abstractmethod
  import torch
label_desc: Use image node_id 2 for calling the NEFTuneHook obj's underlying member method code with example usage: obj.hook_fn(module, input, output) and returns: output