class_docstr: Custom checkpointed function to compute the target score.
class_docstr_tok: ["Custom","checkpointed","function","to","compute","the","target","score","."]
code: def get_target_score(
      i: torch.Tensor,
      w: torch.Tensor,
      target: torch.Tensor,
      full_precision: bool,
      margin: float,
      scale: Optional[float],
  ) -> torch.Tensor:
      tokens, d_model = i.shape
      assert d_model == w.shape[1]
      tw = w.gather(
          dim=0,
          index=target.reshape(target.shape[0], 1).expand(
              target.shape[0], d_model
          ),
      )
      assert tw.shape == (tokens, d_model)
      if scale is not None:
          target_score = F.normalize(i, dim=1) * F.normalize(tw, dim=1)
      else:
          target_score = i * tw
      if full_precision:
          target_score = target_score.float()
      target_score = target_score.sum(
          dim=1
      )  # sum into target scores with shape (tokens,)
      if scale is not None:
          target_score -= margin
          target_score *= scale
      return target_score
  
imports: from typing import Any, Optional, Tuple
  import torch
  from torch import nn
  import torch.distributed
  import torch.nn.functional
label_desc: Use image node_id 1 for calling the TargetScoreFunction obj's underlying member method code with example usage: obj.get_target_score(i, w, target, full_precision, margin, scale) and returns: target_score