Fix multitask finetune

smi-ted/finetune/finetune_classification_multitask.py

@@ -48,6 +48,7 @@ def main(config):
             'MUV-692', 'MUV-712', 'MUV-713', 'MUV-733', 'MUV-737', 'MUV-810',
             'MUV-832', 'MUV-846', 'MUV-852', 'MUV-858', 'MUV-859'
         ]
+    config.n_output = len(targets)
 
     # load dataset
     df_train = pd.read_csv(f"{config.data_root}/train.csv")

smi-ted/finetune/trainers.py

@@ -580,4 +580,4 @@ class TrainerClassifierMultitask(Trainer):
             'specificity': average_sp.item(),
         }
 
-        return preds, running_loss / len(data_loader), metrics
+        return preds.cpu().numpy(), running_loss / len(data_loader), metrics