implement onnxable lasttimestep pooling (#268)

Summary:
Pull Request resolved: #268

add special pooling of the last hidden state when exporting so ATen and extra operators is not introduced

Reviewed By: houseroad

Differential Revision: D13954997

fbshipit-source-id: a33820a75d7b604ef6eb402f2e6898c0ec0890ed

pytext/models/representations/pooling.py

@@ -112,6 +112,14 @@ def __init__(self, config: Module.Config, n_input: int) -> None:
         super().__init__(config)
 
     def forward(self, inputs: torch.Tensor, seq_lengths: torch.Tensor) -> torch.Tensor:
+        # inputs: (bsz, max_len, dim)
+        # seq_lengths: (bsz,)
+
+        if torch._C._get_tracing_state():
+            # if it is exporting, the batch size = 1, so we return the last hidden state
+            # by returning the last dimension to avoid introducing extra operators
+            assert inputs.shape[0] == 1
+            return inputs[:, -1, :]
         bsz, _, dim = inputs.shape
         idx = seq_lengths.unsqueeze(1).expand(bsz, dim).unsqueeze(1)
         return inputs.gather(1, idx - 1).squeeze(1)