optionally checkpoint r_squared metrics when required (pytorch#3083)

Summary:
Pull Request resolved: pytorch#3083

r_squared states get checkpointed when not needed, resulting in loading issues. this bypasses by making those states persistent only when needed, ie: r_squared=True

Reviewed By: zxpmirror1994

Differential Revision: D76162934

fbshipit-source-id: adf76774b8d6331c0a67250eef11060b1bbb419e

Author: iamzainhuda

torchrec/metrics/mse.py

@@ -110,14 +110,14 @@ def __init__(
             torch.zeros(self._n_tasks, dtype=torch.double),
             add_window_state=True,
             dist_reduce_fx="sum",
-            persistent=True,
+            persistent=include_r_squared,
         )
         self._add_state(
             "label_squared_sum",
             torch.zeros(self._n_tasks, dtype=torch.double),
             add_window_state=True,
             dist_reduce_fx="sum",
-            persistent=True,
+            persistent=include_r_squared,
         )
 
     def update(