Fix prediction_writer loop

cellarium-ai · Aug 30, 2024 · 19c69b5 · 19c69b5
1 parent 39e2134
commit 19c69b5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cellarium/ml/callbacks/prediction_writer.py b/cellarium/ml/callbacks/prediction_writer.py
@@ -35,7 +35,7 @@ def write_prediction(
         os.makedirs(output_dir, exist_ok=True)
     df = pd.DataFrame(prediction.cpu())
     if fields is not None:
-        for field_name, field_data in fields:
+        for field_name, field_data in fields.items():
             df.insert(0, field_name, field_data)
     df.insert(0, "db_ids", ids)
     output_path = os.path.join(output_dir, f"batch_{postfix}.csv")