scalar metric fix window count state size (meta-pytorch#3302)

Summary:
Pull Request resolved: meta-pytorch#3302

make window_count size same as other state tensors. no issues with checkpointing as these states are not saved to state dict (persistent=False)

Reviewed By: irobert0126

Differential Revision: D80275353

fbshipit-source-id: 0176fc551a9629480d416eaf4abb896447fa13b2

Author: iamzainhuda

torchrec/metrics/scalar.py

@@ -53,8 +53,10 @@ def update(
 
         states = {
             "labels": labels.mean(dim=-1),
-            "window_count": torch.tensor([1.0]).to(
-                labels.device
+            "window_count": torch.ones(
+                torch.Size([self._n_tasks]),
+                device=labels.device,
+                dtype=torch.double,
             ),  # put window count on the correct device
         }
         for state_name, state_value in states.items():