Fix UnboundLocalError in weights assignment (#2493)

Summary:
Pull Request resolved: #2493

See D64163927 for details. This block is not hit by testing so the typo was not captured.

Reviewed By: Clara85

Differential Revision: D64609231

fbshipit-source-id: b0a4a0ce54dc8ed6d69c35c5a6f9eb77cd70a430

torchrec/distributed/mc_modules.py

@@ -663,7 +663,7 @@ def compute(
                         values=features.values(),
                         lengths=features.lengths(),
                         # TODO: improve this temp solution by passing real weights
-                        weights=torch.tensor(kjt.length_per_key()),
+                        weights=torch.tensor(features.length_per_key()),
                     )
                 }
                 mcm = self._managed_collision_modules[table]