DeepGPR 0.0.18__tar.gz → 0.0.19__tar.gz

{deepgpr-0.0.18 → deepgpr-0.0.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: DeepGPR
-Version: 0.0.18
+Version: 0.0.19
 Summary: PyTorch and CUDA for GPR FWI
 Author-email: Lei Liu <liulei990222@gmail.com>
 License-Expression: MIT

{deepgpr-0.0.18 → deepgpr-0.0.19}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "DeepGPR"
-version = "0.0.18"
+version = "0.0.19"
 authors = [
     { name = "Lei Liu", email = "liulei990222@gmail.com" }
 ]

{deepgpr-0.0.18 → deepgpr-0.0.19}/src/DeepGPR/compute2.py

@@ -445,23 +445,32 @@ def _save_forward_wavefield(wavefield, directory, run_time):
 
 
 def _begin_native_call(device):
-    """Order the native default stream after PyTorch's current CUDA stream."""
+    """Select the tensor device and order the native CUDA default stream."""
     if device.type != "cuda":
         return None
 
+    previous_device = torch.cuda.current_device()
+    target_device = previous_device if device.index is None else device.index
+    restore_device = None
+    if target_device != previous_device:
+        torch.cuda.set_device(target_device)
+        restore_device = previous_device
+
     current_stream = torch.cuda.current_stream(device)
     default_stream = torch.cuda.default_stream(device)
     if current_stream.cuda_stream != default_stream.cuda_stream:
         default_stream.wait_stream(current_stream)
-        return current_stream, default_stream
-    return None
+    return current_stream, default_stream, restore_device
 
 
-def _end_native_call(streams):
-    """Make PyTorch's current stream wait for an asynchronous native call."""
-    if streams is not None:
-        current_stream, default_stream = streams
-        current_stream.wait_stream(default_stream)
+def _end_native_call(native_state):
+    """Order native results before PyTorch and restore the caller's device."""
+    if native_state is not None:
+        current_stream, default_stream, restore_device = native_state
+        if current_stream.cuda_stream != default_stream.cuda_stream:
+            current_stream.wait_stream(default_stream)
+        if restore_device is not None:
+            torch.cuda.set_device(restore_device)
 
 
 def _check_nonzero_source_created_fields(c_lib, source_amplitudes, *fields):

{deepgpr-0.0.18 → deepgpr-0.0.19}/src/DeepGPR.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: DeepGPR
-Version: 0.0.18
+Version: 0.0.19
 Summary: PyTorch and CUDA for GPR FWI
 Author-email: Lei Liu <liulei990222@gmail.com>
 License-Expression: MIT

{deepgpr-0.0.18 → deepgpr-0.0.19}/tests/test_discrete_adjoint.py

@@ -279,6 +279,51 @@ class DiscreteAdjointTests(unittest.TestCase):
                     2.0e-4,
                 )
 
+    @unittest.skipUnless(
+        torch.cuda.is_available() and torch.cuda.device_count() >= 2,
+        "At least two CUDA devices are required",
+    )
+    def test_cuda_explicit_device_selects_matching_native_runtime(self):
+        original_device = torch.cuda.current_device()
+        target_index = (original_device + 1) % torch.cuda.device_count()
+        target_device = torch.device(f"cuda:{target_index}")
+        torch.cuda.set_device(original_device)
+
+        try:
+            shape = (10, 11, 9)
+            nt = 8
+            eps_r = torch.full(shape, 4.0, device=target_device)
+            sigma = torch.zeros_like(eps_r)
+            source = torch.zeros((1, nt, 1), device=target_device)
+            source[0, 1, 0] = 1.0
+            source_location = torch.tensor(
+                [[[5, 5, 4]]], dtype=torch.int32, device=target_device
+            )
+            receiver_location = torch.tensor(
+                [[[5, 7, 4]]], dtype=torch.int32, device=target_device
+            )
+
+            receiver = DeepGPR.compute(
+                device=target_device,
+                dx=0.02,
+                dt=1.0e-11,
+                source_amplitudes=source,
+                source_location=source_location,
+                receiver_location=receiver_location,
+                eps_r=eps_r,
+                sigma=sigma,
+                pmlthick=2,
+                fdtd_order=2,
+                mode=3,
+            )[-1]
+            torch.cuda.synchronize(target_device)
+
+            self.assertEqual(receiver.device, target_device)
+            self.assertTrue(torch.isfinite(receiver).all().item())
+            self.assertEqual(torch.cuda.current_device(), original_device)
+        finally:
+            torch.cuda.set_device(original_device)
+
     def test_source_waveform_gradient_matches_finite_difference(self):
         nx, ny, nt = 12, 14, 60
         eps_r = torch.full((nx, ny), 4.0)