From ead7a4572a530f750acf2b7393f3b18f38df58d9 Mon Sep 17 00:00:00 2001
From: Louis Abraham <louisabraham@users.noreply.github.com>
Date: Mon, 8 Jan 2024 14:44:20 +0100
Subject: [PATCH] fix #52

---
 lassonet/cox.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/lassonet/cox.py b/lassonet/cox.py
index facc338..562975f 100644
--- 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()