Fixes issue where tensors in waveform generator not getting built on …

…correct device (#199) * fix device issue in waveform generator * update version to 0.7.1
ML4GW · Feb 7, 2025 · a55d194 · a55d194
1 parent b9895af
commit a55d194
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/ml4gw/waveforms/generator.py b/ml4gw/waveforms/generator.py
@@ -224,7 +224,7 @@ def generate_conditioned_fd_waveform(
         k1s = torch.round(f_min / df)
 
         num_freqs = frequencies.size(0)
-        frequency_indices = torch.arange(num_freqs)
+        frequency_indices = torch.arange(num_freqs, device=device)
         taper_mask = frequency_indices <= k1s[:, None]
         taper_mask &= frequency_indices >= k0s[:, None]
 
@@ -253,7 +253,7 @@ def generate_conditioned_fd_waveform(
         # that will translate the coalescense time such that it is `right_pad`
         # seconds from the right edge of the window
         tshift = round(self.right_pad * self.sample_rate) / self.sample_rate
-        kvals = torch.arange(num_freqs)
+        kvals = torch.arange(num_freqs, device=device)
         phase_shift = torch.exp(1j * 2 * torch.pi * df * tshift * kvals)
 
         hc_spectrum *= phase_shift

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ml4gw"
-version = "0.7.0"
+version = "0.7.1"
 description = "Tools for training torch models on gravitational wave data"
 readme = "README.md"
 authors = [