feat: pass batch size to predict in dynamic deephit (#153)

vanderschaarlab · Mar 22, 2023 · e45ced0 · e45ced0
1 parent f117206
commit e45ced0
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/synthcity/plugins/core/models/time_series_survival/ts_surv_dynamic_deephit.py b/src/synthcity/plugins/core/models/time_series_survival/ts_surv_dynamic_deephit.py
@@ -141,6 +141,7 @@ def predict(
         temporal: Union[np.ndarray, List],
         observation_times: Union[np.ndarray, List],
         time_horizons: List,
+        batch_size: Optional[int] = 100,
     ) -> np.ndarray:
         "Predict risk"
         static = np.asarray(static)
@@ -150,7 +151,8 @@ def predict(
         data = self._merge_data(static, temporal, observation_times)
 
         return pd.DataFrame(
-            self.model.predict_risk(data, time_horizons), columns=time_horizons
+            self.model.predict_risk(data, time_horizons, bs=batch_size),
+            columns=time_horizons,
         )
 
     @validate_arguments(config=dict(arbitrary_types_allowed=True))
@@ -221,7 +223,6 @@ def __init__(
         clipping_value: int = 1,
         output_type: str = "MLP",
     ) -> None:
-
         self.split = split
         self.split_time = None