class_docstr: Keep track of how many tokens we've seen.
      
class_docstr_tok: ["Keep","track","of","how","many","tokens","we","'ve","seen","."]
code: def __init__(self, **kwargs: Dict[str, Any]):
      super().__init__(**kwargs)
      self.add_state(
          "count",
          default=torch.tensor(0, dtype=torch.int64),
          dist_reduce_fx="sum",
          persistent=True,
      )
  
file_path: flash-attention/training/src/metrics/num_tokens.py
imports: from typing import Any, Dict, Optional
  import torch
  from torch import Tensor
  from torchmetrics import Metric
label_desc: Use image node_id 1 to create a new NumTokens object from inherited base classes: Metric with example: obj = NumTokens()