fix batch inputs bug

ML4GW · Jan 9, 2025 · 13f61f6 · 13f61f6
1 parent 368d0dd
commit 13f61f6
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/ml4gw/waveforms/cbc/phenom_p.py b/ml4gw/waveforms/cbc/phenom_p.py
@@ -379,10 +379,10 @@ def PhenomPOneFrequency(
         diffRDphase = (diff[:, 1:] + diff[:, :-1]) / (
             2 * delta_fRds.unsqueeze(1)
         )
+        # reshape x to have same shape as diffRDphase
+        x = x[1:-1].unsqueeze(0).expand(diffRDphase.shape)
         # interpolate at x = 1, as thats the same as f = fRD
-        diffRDphase = -self.interpolate(
-            torch.tensor([1]), x[1:-1], diffRDphase
-        )
+        diffRDphase = -self.interpolate(torch.tensor([1]), x, diffRDphase)
         return hPhenom, diffRDphase
 
     # Utility functions