Merge branch 'master' of https://github.com/ilemhadri/lassonet

lasso-net · Jan 12, 2024 · f812926 · f812926
2 parents 3b0d1b4 + ead7a45
commit f812926
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/lassonet/cox.py b/lassonet/cox.py
@@ -31,7 +31,9 @@ def forward(self, log_h, y):
         events = events[idx]
 
         event_ind = events.nonzero().flatten()
-
+        if event_ind.nelement() == 0:
+            return torch.tensor(0.0)
+
         # numerator
         log_num = log_h[event_ind].mean()