class_docstr: Adds PyTorch specific functions.
class_docstr_tok: ["Adds","PyTorch","specific","functions","."]
docstr: Perform one training step over data in PyTorch models.
  
  Returns
  -------
  (LP, LA) : tuple of (float, float)
      predictor loss and adversary loss.
docstr_tok: ["Perform","one","training","step","over","data","in","PyTorch","models",".","Returns","--","--","--","-","(","LP",",","LA",")",":","tuple","of","(","float",",","float",")","predictor","loss","and","adversary","loss","."]
code: def train_step(self, X, Y, A):
      """
      Perform one training step over data in PyTorch models.
  
      Returns
      -------
      (LP, LA) : tuple of (float, float)
          predictor loss and adversary loss.
      """
      self.predictor_model.train()
      self.adversary_model.train()
  
      # Clear gradient
      self.predictor_optimizer.zero_grad()
      self.adversary_optimizer.zero_grad()
  
      Y_hat = self.predictor_model(X)
      LP = self.predictor_loss(Y_hat, Y)
      LP.backward(
          retain_graph=True
      )  # Check what this does at some point in time
  
      dW_LP = [
          torch.clone(p.grad.detach())
          for p in self.predictor_model.parameters()
      ]
  
      self.predictor_optimizer.zero_grad()
      self.adversary_optimizer.zero_grad()
  
      # For equalized odds
      if self.base.pass_y_:
          Y_hat = torch.cat((Y_hat, Y), dim=1)
  
      A_hat = self.adversary_model(Y_hat)
      LA = self.adversary_loss(A_hat, A)
      LA.backward()
  
      dW_LA = [
          torch.clone(p.grad.detach())
          for p in self.predictor_model.parameters()
      ]
  
      for i, p in enumerate(self.predictor_model.parameters()):
          # Normalize dW_LA
          unit_dW_LA = dW_LA[i] / (
              torch.norm(dW_LA[i]) + torch.finfo(float).tiny
          )
          # Project
          proj = torch.sum(torch.inner(unit_dW_LA, dW_LP[i]))
          # Calculate dW
          p.grad = (
              dW_LP[i]
              - (proj * unit_dW_LA)
              - (self.base.alpha * dW_LA[i])
          )
  
      self.predictor_optimizer.step()
      self.adversary_optimizer.step()
  
      return (LP.item(), LA.item())
  
file_path: fairlearn/fairlearn/adversarial/_pytorch_engine.py
imports: from ._backend_engine import BackendEngine
  from ._constants import _MODEL_UNRECOGNIZED_STR, _MODEL_UNRECOGNIZED_ITEM
label_desc: Use image node_id 5 for calling the PytorchEngine obj's underlying member method code with example usage: obj.train_step(X, Y, A) and returns: LP,  LA