code: def _get_states(
      labels: torch.Tensor,
      predictions: torch.Tensor,
      weights: torch.Tensor,
  ) -> Dict[str, torch.Tensor]:
      cross_entropy = compute_cross_entropy(
          labels, predictions, weights, TestNEMetric.eta
      )
      cross_entropy_sum = torch.sum(cross_entropy)
      weighted_num_samples = torch.sum(weights)
      pos_labels = torch.sum(weights * labels)
      neg_labels = torch.sum(weights * (1.0 - labels))
      return {
          "cross_entropy_sum": cross_entropy_sum,
          "weighted_num_samples": weighted_num_samples,
          "pos_labels": pos_labels,
          "neg_labels": neg_labels,
          "num_samples": torch.tensor(labels.size()).long(),
      }
  
imports: import unittest
  from functools import partial, update_wrapper
  from typing import Callable, Dict, Type
  import torch
  from torchrec.metrics.ne import compute_cross_entropy, compute_logloss, compute_ne, NEMetric
  from torchrec.metrics.rec_metric import RecComputeMode, RecMetric
  from torchrec.metrics.test_utils import metric_test_helper, rec_metric_value_test_launcher, TestMetric
label_desc: Use image node_id 1 for calling the TestNEMetric obj's underlying member method code with example usage: obj._get_states(labels, predictions, weights) and returns: dict