Fix UnboundLocalError in weights assignment (#2493)

Ziliang Zhao · facebook-github-bot · commit 1d8824b1707f · 2024-10-18T11:54:33.000-07:00
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
diff --git a/torchrec/distributed/mc_modules.py b/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]

Original file line number	Diff line number	Diff line change
`@@ -663,7 +663,7 @@ def compute(`
`663`	`663`	`values=features.values(),`
`664`	`664`	`lengths=features.lengths(),`
`665`	`665`	`# TODO: improve this temp solution by passing real weights`
`666`		`- weights=torch.tensor(kjt.length_per_key()),`
	`666`	`+ weights=torch.tensor(features.length_per_key()),`
`667`	`667`	`)`
`668`	`668`	`}`
`669`	`669`	`mcm = self._managed_collision_modules[table]`