class_docstr: ## Gradient Penalty
class_docstr_tok: ["#","#","Gradient","Penalty"]
docstr: * `x` is $x \sim \mathbb{P}_r$
  * `f` is $D(x)$
  
  $\hat{x} \leftarrow x$
  since we set $\epsilon = 1$ for this implementation.
docstr_tok: ["*","`","x","`","is","$","x","\\sim","\\mathbb","{","P","}","_r","$","*","`","f","`","is","$","D","(","x",")","$","$","\\hat","{","x","}","\\leftarrow","x","$","since","we","set","$","\\epsilon","=","1","$","for","this","implementation","."]
code: def forward(self, x: torch.Tensor, f: torch.Tensor):
      """
      * `x` is $x \sim \mathbb{P}_r$
      * `f` is $D(x)$
  
      $\hat{x} \leftarrow x$
      since we set $\epsilon = 1$ for this implementation.
      """
  
      # Get batch size
      batch_size = x.shape[0]
  
      # Calculate gradients of $D(x)$ with respect to $x$.
      # `grad_outputs` is set to ones since we want the gradients of $D(x)$,
      # and we need to create and retain graph since we have to compute gradients
      # with respect to weight on this loss.
      gradients, *_ = torch.autograd.grad(
          outputs=f,
          inputs=x,
          grad_outputs=f.new_ones(f.shape),
          create_graph=True,
      )
  
      # Reshape gradients to calculate the norm
      gradients = gradients.reshape(batch_size, -1)
      # Calculate the norm $\Vert \nabla_{\hat{x}} D(\hat{x}) \Vert_2$
      norm = gradients.norm(2, dim=-1)
      # Return the loss $\big(\Vert \nabla_{\hat{x}} D(\hat{x}) \Vert_2 - 1\big)^2$
      return torch.mean((norm - 1) ** 2)
  
imports: import torch
  import torch.autograd
  from labml_helpers.module import Module
label_desc: Use image node_id 1 for calling the GradientPenalty obj's underlying member method code with example usage: obj.forward(x, f) and returns: torch