bug fix

AI4EPS · Sep 28, 2024 · 737551f · 737551f
1 parent fa2585f
commit 737551f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/adloc/adloc.py b/adloc/adloc.py
@@ -225,6 +225,7 @@ def __init__(
 
         self.event_loc.weight.requires_grad = True
         self.event_time.weight.requires_grad = False
+        # self.event_time.weight.requires_grad = True
 
     def calc_time(self, event_loc, station_loc, phase_type):
         if self.eikonal is None:
@@ -313,6 +314,8 @@ def forward(
                 phase_time_ = phase_time[phase_type == type]
                 loss += torch.sum(F.huber_loss(t_, phase_time_, reduction="none") * phase_weight_)
 
+        if loss == 0.0:
+            return None
         return {"phase_time": pred_time, "loss": loss}
 
 

diff --git a/adloc/data.py b/adloc/data.py
@@ -231,6 +231,10 @@ def __len__(self):
     def __getitem__(self, i):
 
         idx = self.idx_batch[i][self.valid_index[self.idx_batch[i]]]
+
+        if len(idx) == 0:
+            return None  # skip empty batch
+
         idx1_eve = self.pairs["event_index1"][idx]
         idx2_eve = self.pairs["event_index2"][idx]
         idx_eve = np.stack([idx1_eve, idx2_eve], axis=1)