GAUC fix (pytorch#2619)

Summary: Pull Request resolved: pytorch#2619 fix gauc multiple returns Reviewed By: yunjiangster, arsatis Differential Revision: D66853683 fbshipit-source-id: cf5f47a22086b2d225dc38b6dd3d90b25e53cc93
iamzainhuda · Dec 6, 2024 · f450c59 · f450c59
1 parent 77640f8
commit f450c59
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchrec/metrics/gauc.py b/torchrec/metrics/gauc.py
@@ -209,8 +209,8 @@ def _compute(self) -> List[MetricComputationReport]:
                 name=MetricName.GAUC_NUM_SAMPLES,
                 metric_prefix=MetricPrefix.LIFETIME,
                 value=compute_window_auc(
-                    self.get_window_state("auc_sum"),
-                    self.get_window_state("num_samples"),
+                    cast(torch.Tensor, self.auc_sum),
+                    cast(torch.Tensor, self.num_samples),
                 )["num_samples"],
             ),
             MetricComputationReport(