class_docstr: Computes the confusion matrix and is based on `FCN <https://github.com/shelhamer/fcn.berkeleyvision.org/blob/master/score.py>`_
code: def update(
      self,
      prediction: Union[Tensor, Dict],
      target: Union[Tensor, Dict],
      extras: Dict[str, Any] = {},
      batch_size: Optional[int] = 1,
  ):
      if (
          isinstance(prediction, dict)
          and self.prediction_key in prediction
      ):
          prediction = prediction[self.prediction_key]
  
      if isinstance(prediction, dict) or isinstance(prediction, dict):
          raise NotImplementedError(
              "ConfusionMatrix does not currently support Dict predictions or targets"
          )
  
      n_classes = prediction.shape[1]
  
      if self.confusion_mat is None:
          self.confusion_mat = torch.zeros(
              (n_classes, n_classes),
              dtype=torch.int64,
              device=target.device,
          )
  
      with torch.no_grad():
          prediction = prediction.argmax(1).flatten()
          target = target.flatten()
          k = (target >= 0) & (target < n_classes)
          inds = n_classes * target[k].to(torch.int64) + prediction[k]
          cnts = torch.bincount(inds, minlength=n_classes**2).reshape(
              n_classes, n_classes
          )
          if self.is_distributed:
              cnts = reduce_tensor_sum(cnts)
          self.confusion_mat += cnts
  
file_path: ml-cvnets/metrics/confusion_mat.py
imports: from numbers import Number
  from typing import Any, Dict, List, Optional, Union
  import torch
  from torch import Tensor
  from metrics import METRICS_REGISTRY
  from metrics.metric_base import BaseMetric
  from utils.tensor_utils import reduce_tensor_sum
label_desc: Use image node_id 2 for calling the ConfusionMatrix obj's underlying member method code with example usage: obj.update(prediction, target, extras, batch_size) without return types