ffn: fix error for loss kwargs

Author: aditya0by0

chebai/models/ffn.py

@@ -37,7 +37,7 @@ def _get_prediction_and_labels(self, data, labels, model_output):
         loss_kwargs = data.get("loss_kwargs", dict())
         if "non_null_labels" in loss_kwargs:
             n = loss_kwargs["non_null_labels"]
-            d = data[n]
+            d = d[n]
         return torch.sigmoid(d), labels.int() if labels is not None else None
 
     def _process_for_loss(