TextFieldTensor in multitask models (#5331)

* update `make_inputs_for_task` to support TextFieldTensors

* Changelog

* Formatting

Co-authored-by: Amit Parekh <7276308+amitkparekh@users.noreply.github.com>

CHANGELOG.md

@@ -32,6 +32,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `TransformerTextField` can now take tensors of shape `(1, n)` like the tensors produced from a HuggingFace tokenizer.
 - `tqdm` lock is now set inside `MultiProcessDataLoading` when new workers are spawned to avoid contention when writing output.
 - `ConfigurationError` is now pickleable.
+- Multitask models now support `TextFieldTensor` in heads, not just in the backbone
 
 ### Changed
 

allennlp/models/multitask.py

@@ -5,7 +5,7 @@
 from overrides import overrides
 import torch
 
-from allennlp.data import Vocabulary
+from allennlp.data import Vocabulary, TextFieldTensors
 from allennlp.modules import Backbone
 from allennlp.models.model import Model
 from allennlp.models.heads import Head
@@ -111,7 +111,16 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:  # type: ignore
             task: torch.LongTensor(indices) for task, indices in task_indices_just_for_mypy.items()
         }
 
-        def make_inputs_for_task(task: str, whole_batch_input: Union[torch.Tensor, List]):
+        def make_inputs_for_task(
+            task: str, whole_batch_input: Union[torch.Tensor, TextFieldTensors, List]
+        ):
+            if isinstance(whole_batch_input, dict):
+                for k1, v1 in whole_batch_input.items():
+                    for k2, v2 in v1.items():
+                        whole_batch_input[k1][k2] = make_inputs_for_task(task, v2)
+
+                return whole_batch_input
+
             if isinstance(whole_batch_input, torch.Tensor):
                 task_indices[task] = task_indices[task].to(whole_batch_input.device)
                 return torch.index_select(whole_batch_input, 0, task_indices[task])