Skip to content
This repository was archived by the owner on Nov 22, 2022. It is now read-only.

Commit 4be819b

Browse files
shreydesaifacebook-github-bot
authored andcommitted
Making seq2seq_model more torchscript-friendly (#1459)
Summary: Pull Request resolved: #1459 seq2seq_model.py changes: - removes dependence on _ when unpacking the tensor dicts - explicit None checks on dict feats Differential Revision: D23673319 fbshipit-source-id: dd0f8e164b7368a6eef7fcf02a1a15a3c078f33e
1 parent 5b5b6eb commit 4be819b

File tree

1 file changed

+5
-5
lines changed

1 file changed

+5
-5
lines changed

pytext/models/seq_models/seq2seq_model.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -107,8 +107,8 @@ def arrange_model_inputs(
107107
torch.Tensor,
108108
torch.Tensor,
109109
]:
110-
src_tokens, src_lengths, _ = tensor_dict["src_seq_tokens"]
111-
trg_tokens, trg_lengths, _ = tensor_dict["trg_seq_tokens"]
110+
src_tokens, src_lengths = tensor_dict["src_seq_tokens"]
111+
trg_tokens, trg_lengths = tensor_dict["trg_seq_tokens"]
112112

113113
def _shift_target(in_sequences, seq_lens, eos_idx, pad_idx):
114114
shifted_sequence = GetTensor(
@@ -136,7 +136,7 @@ def _shift_target(in_sequences, seq_lens, eos_idx, pad_idx):
136136
)
137137

138138
def arrange_targets(self, tensor_dict):
139-
trg_tokens, trg_lengths, _ = tensor_dict["trg_seq_tokens"]
139+
trg_tokens, trg_lengths = tensor_dict["trg_seq_tokens"]
140140
return (trg_tokens, trg_lengths)
141141

142142
def __init__(
@@ -196,7 +196,7 @@ def forward(
196196
):
197197
additional_features: List[List[torch.Tensor]] = []
198198

199-
if dict_feats:
199+
if dict_feats is not None:
200200
additional_features.append(list(dict_feats))
201201

202202
if contextual_token_embedding is not None:
@@ -206,7 +206,7 @@ def forward(
206206
src_tokens, additional_features, src_lengths, trg_tokens
207207
)
208208

209-
if dict_feats:
209+
if dict_feats is not None:
210210
(
211211
output_dict["dict_tokens"],
212212
output_dict["dict_weights"],

0 commit comments

Comments
 (0)