Skip to content

Commit bd96e8f

Browse files
PaulZhang12facebook-github-bot
authored andcommitted
Fix CPU Unit Test CI on OSS (#1777)
Summary: Pull Request resolved: #1777 {F1467318350} CPU Unit Test CI for OSS failing on call to PT2 ExportedProgram: https://github.com/pytorch/torchrec/actions/runs/8233905334/job/22514441384. OSS CI expects .module() instead of directly calling forward Reviewed By: joshuadeng, gnahzg Differential Revision: D54751106 fbshipit-source-id: d1c614ec593edaf31ca894c36d2430e06f779f76
1 parent 0553999 commit bd96e8f

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

torchrec/distributed/tests/test_pt2.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ def _test_kjt_input_module(
132132

133133
if test_pt2_ir_export:
134134
pt2_ir = torch.export.export(EM, inputs, {}, strict=False)
135-
pt2_ir_output = pt2_ir(*inputs)
135+
pt2_ir_output = pt2_ir.module()(*inputs)
136136
assert_close(eager_output, pt2_ir_output)
137137

138138
def test_kjt_split(self) -> None:

0 commit comments

Comments
 (0)