mark metricmodule methods as experimental (#2983)

Summary:

tsia

Reviewed By: xunnanxu, kausv

Differential Revision: D75014609

Author: iamzainhuda

torchrec/metrics/metric_module.py

@@ -68,6 +68,7 @@
 from torchrec.metrics.unweighted_ne import UnweightedNEMetric
 from torchrec.metrics.weighted_avg import WeightedAvgMetric
 from torchrec.metrics.xauc import XAUCMetric
+from torchrec.utils.experimental import experimental
 
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -394,6 +395,7 @@ def _get_metric_states(
 
         return state_aggregated
 
+    @experimental
     def get_pre_compute_states(
         self, pg: Optional[Union[dist.ProcessGroup, DeviceMesh]] = None
     ) -> Dict[str, Dict[str, Dict[str, Union[torch.Tensor, List[torch.Tensor]]]]]:
@@ -442,6 +444,7 @@ def get_pre_compute_states(
 
         return aggregated_states
 
+    @experimental
     def load_pre_compute_states(
         self,
         source: Dict[

torchrec/metrics/tests/test_metric_module.py

@@ -643,7 +643,7 @@ def metric_module_gather_state(
             metric_module.update(test_batch)
 
         computed_value = metric_module.compute()
-        states = metric_module.get_pre_compute_states(pg=ctx.pg)  # pyre-ignore[6]
+        states = metric_module.get_pre_compute_states(pg=ctx.pg)
 
         torch.distributed.barrier(ctx.pg)
         # Compare to computing metrics on metric module that loads from pre_compute_states