Fix: Test failing when using DeepXDE

tests/pde/deepxde_example.py

@@ -1,5 +1,4 @@
 import torch
-import deepxde as dde
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -12,6 +11,10 @@ def deepxde_example():
     https://deepxde.readthedocs.io/en/latest/demos/operator/poisson.1d.pideeponet.html
     """
 
+    # deepxde sets the device context internally, which can conflict with the testing setup, when dealing with different
+    # devices (i.e. GPU and CPU). To ensure that the correct device is set the dependency is isolated.
+    import deepxde as dde  # noqa
+
     # Poisson equation: -u_xx = f
     def equation(x, y, f):
         dy_xx = dde.grad.hessian(y, x)

tests/pde/test_pideeponet.py

@@ -1,6 +1,5 @@
 import pytest
 import torch
-import deepxde as dde
 import matplotlib.pyplot as plt
 import numpy as np
 import continuiti as cti
@@ -15,6 +14,10 @@ def test_pideeponet():
     https://deepxde.readthedocs.io/en/latest/demos/operator/poisson.1d.pideeponet.html
     """
 
+    # deepxde sets the device context internally, which can conflict with the testing setup, when dealing with different
+    # devices (i.e. GPU and CPU). To ensure that the correct device is set the dependency is isolated.
+    import deepxde as dde  # noqa
+
     # Poisson equation: -v_xx = f
     mse = torch.nn.MSELoss()