Change multi risk for irregular length (#32)

autonlab · Jan 21, 2021 · eb864a3 · eb864a3
1 parent 6d6016b
commit eb864a3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dsm/dsm_api.py b/dsm/dsm_api.py
@@ -110,7 +110,7 @@ def fit(self, x, t, e, vsize=0.15,
     x_train, t_train, e_train, x_val, t_val, e_val = processed_data
 
     inputdim = x_train.shape[-1]
-    maxrisk = int(e_train.max())
+    maxrisk = int(np.nanmax(e_train.cpu().numpy()))
     model = self._gen_torch_model(inputdim, optimizer, risks=maxrisk)
     model, _ = train_dsm(model,
                          x_train, t_train, e_train,