fix task index

sileod · Apr 24, 2023 · e5908d9 · e5908d9
1 parent 083a9b3
commit e5908d9
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 7 deletions.
diff --git a/src/tasknet/models.py b/src/tasknet/models.py
@@ -266,12 +266,8 @@ def __init__(self, tasks):
     def __call__(
         self, features: List[Union[InputDataClass, Dict]]
     ) -> Dict[str, torch.Tensor]:
-        #try:
-        task_index = features[0]["task"].flatten()[0].item()
-        #except:
-        #    print("features:",features)
-        #   task_index = features[-1]["task"].flatten()[0].item()
-
+
+        task_index = features[0]["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])

diff --git a/src/tasknet/tasks.py b/src/tasknet/tasks.py
@@ -45,7 +45,9 @@ def sample_dataset(dataset,n=10000, n_eval=1000, oversampling=None):
 
 def get_len(outputs):
     try:
-        return len(outputs[fc.first(outputs)])
+        batch_length=len(outputs[fc.first(outputs)])
+        assert batch_length
+        return batch_length
     except:
         return 1