Skip to content

Commit fc7e83d

Browse files
wilson100hongfacebook-github-bot
authored andcommitted
register_buffer for DLRM.triu_indices to make it exportable (#1670)
Summary: Pull Request resolved: #1670 we need to register buffer for `triu_indices` to make DLRM "exportable" -- thus to generate PT2 IR. Reviewed By: joshuadeng Differential Revision: D53233580 fbshipit-source-id: 451c70ea5e3433756b3a351f74e9eedc99cfca0b
1 parent 584bb44 commit fc7e83d

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

torchrec/models/dlrm.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -183,8 +183,10 @@ class InteractionArch(nn.Module):
183183
def __init__(self, num_sparse_features: int) -> None:
184184
super().__init__()
185185
self.F: int = num_sparse_features
186-
self.triu_indices: torch.Tensor = torch.triu_indices(
187-
self.F + 1, self.F + 1, offset=1
186+
self.register_buffer(
187+
"triu_indices",
188+
torch.triu_indices(self.F + 1, self.F + 1, offset=1),
189+
persistent=False,
188190
)
189191

190192
def forward(

0 commit comments

Comments
 (0)