collator

sileod · Feb 27, 2023 · 5431ecb · 5431ecb
1 parent 4b26ff5
commit 5431ecb
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/tasknet/models.py b/src/tasknet/models.py
@@ -213,7 +213,12 @@ def __init__(self, tasks):
     def __call__(
         self, features: List[Union[InputDataClass, Dict]]
     ) -> Dict[str, torch.Tensor]:
-        task_index = features[0]["task"].flatten()[0].item()
+        try:
+            task_index = features[0]["task"].flatten()[0].item()
+        except:
+            print("features:",features)
+            task_index = features[-1]["task"].flatten()[0].item()
+
         features = [{k:v for k,v in x.items() if k!='task'} for x in features]
         collated = self.tasks[task_index].data_collator.__call__(features)
         collated['task']=torch.tensor([task_index])