wrapping to jt_dict (#1306)

Leon Gao · facebook-github-bot · commit 0db5ec5eb518 · 2023-08-04T14:07:31.000-07:00
Summary: Pull Request resolved: #1306 * a simple fx.wrap Differential Revision: D48080861 Privacy Context Container: L1138451 fbshipit-source-id: 90f2269d494120f0242366ba4e8d09fa4790e5d9
diff --git a/torchrec/sparse/jagged_tensor.py b/torchrec/sparse/jagged_tensor.py
@@ -710,6 +710,7 @@ def forward(
         )
 
 
+@torch.fx.wrap
 def _maybe_compute_kjt_to_jt_dict(
     stride: int,
     keys: List[str],

Original file line number	Diff line number	Diff line change
`@@ -710,6 +710,7 @@ def forward(`
`710`	`710`	`)`
`711`	`711`
`712`	`712`
	`713`	`+@torch.fx.wrap`
`713`	`714`	`def _maybe_compute_kjt_to_jt_dict(`
`714`	`715`	`stride: int,`
`715`	`716`	`keys: List[str],`