Enhance mortality label check for string flags (#551)

The current check will set all labels in MIMIC-IV to 0.

Author: v1xerunt

pyhealth/tasks/mortality_prediction.py

@@ -207,7 +207,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]:
             next_admission = admissions[i + 1]
 
             # Check discharge status for mortality label - more robust handling
-            if next_admission.hospital_expire_flag not in [0, 1]:
+            if next_admission.hospital_expire_flag not in [0, 1, "0", "1"]:
                 mortality_label = 0
             else:
                 mortality_label = int(next_admission.hospital_expire_flag)
@@ -355,7 +355,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]:
             next_admission = admissions[i + 1]
 
             # Check discharge status for mortality label
-            if next_admission.hospital_expire_flag not in [0, 1]:
+            if next_admission.hospital_expire_flag not in [0, 1, "0", "1"]:
                 mortality_label = 0
             else:
                 mortality_label = int(next_admission.hospital_expire_flag)