Merge pull request #81 from vkazei/master

fixes device mapping in location_interpolation.py
ar4 · Oct 25, 2024 · 00798f8 · 00798f8
2 parents 5386203 + 762e451
commit 00798f8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/deepwave/location_interpolation.py b/src/deepwave/location_interpolation.py
@@ -331,7 +331,7 @@ def receiver(self,
             for i in range(n_per_shot):
                 out[shotidx, i, :] = (
                     amplitudes[shotidx, self.idxs[hicks_shotidx][i]] *
-                    (self.weights[hicks_shotidx][i][0].reshape(-1, 1) *
-                     self.weights[hicks_shotidx][i][1].reshape(
+                    (self.weights[hicks_shotidx][i][0].to(amplitudes.device).reshape(-1, 1) *
+                     self.weights[hicks_shotidx][i][1].to(amplitudes.device).reshape(
                          1, -1)).reshape(-1)[..., None]).sum(dim=0)
         return out