def update(self, preds: torch.Tensor, target: torch.Tensor):
        """Updates the internal state with results from a new batch.

        Args:
            preds (~torch.Tensor): The predictions from the model, a Tensor of logits.
            target (~torch.Tensor): A Tensor of ground-truth token values.
        """
        # Convert logits to predicted token indices
        preds = torch.argmax(preds, dim=-1)

        # Create mask for non-ignored tokens
        mask = (target != self.ignore_index)
        masked_target = target[mask]
        masked_preds = preds[mask]

        assert isinstance(self.correct_tokens, torch.Tensor)
        assert isinstance(self.total_tokens, torch.Tensor)

        # Update correct and total counts
        self.correct_tokens += torch.sum(masked_preds == masked_target)
        self.total_tokens += masked_target.numel()
Updates the internal state with results from a new batch.

        Args:
            preds (~torch.Tensor): The predictions from the model, a Tensor of logits.
            target (~torch.Tensor): A Tensor of ground-truth token values.
        
llmfoundry/metrics/token_acc.py