RuntimeError when importing deepxde in unrelated test

[JakobEliasWagner]

I encountered an issue while running tests in my project. Specifically, the error arises in tests/test_loader.py when executing the test_loader function. This error only manifests when importing DeepXDE in an unrelated test. The test setup involves using a DataLoader with a RandomDataset and importing DeepXDE in a second unrelated test. The detailed error traceback and minimal example to recreate issue is provided below.

https://github.com/JakobEliasWagner/deepxde-pytest-generator-device/tree/main

Steps to Reproduce

1. Set up a test environment with the following configurations:
   a. Python 3.12
   b. PyTorch (version 2.3.1)
   c. deepxde (version 1.11.1)
2. Create a RandomDataset and use it with a DataLoader in a test function.
3. Import DeepXDE in an unrelated file that is parsed during the pytest build stage (does not have to be part of any test).
4. Run tests with CUDA support enabled.

Expected Behavior

DeepXDE should not modify any devices or generators in the background as this would complicate using it as a dependency.

Actual Behavior

PyTorch Dataloader does not have the correct generator set. Manually setting the generator does fix the issue. Additionally removing the DeepXDE import statement does resolve the issues as well.

Request for Assistance:

I would appreciate any guidance on resolving this issue, whether it be a configuration change, workaround, or a fix in the library itself. Thank you!

# some_test.py
import torch
from torch.utils.data import DataLoader, Dataset


class RandomDataset(Dataset):
    def __init__(self, n: int):
        super().__init__()
        self.n = n
        self.data = torch.rand(n)

    def __len__(self):
        return self.n

    def __getitem__(self, idx):
        return self.data[idx]


def test_loader():
    dataset = RandomDataset(100)
    loader = DataLoader(dataset, batch_size=10, shuffle=True)

    for x in loader:
        y = x ** 2
        assert isinstance(y, torch.Tensor)

# some_test_with_deepxde_import.py
import deepxde as dde

# trace
tests/test_loader.py:17 (test_loader)
def test_loader():
        dataset = RandomDataset(100)
        loader = DataLoader(dataset, batch_size=10, shuffle=True)
    
>       for x in loader:

tests/test_loader.py:22: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
.venv/lib/python3.12/site-packages/torch/utils/data/dataloader.py:631: in __next__
    data = self._next_data()
.venv/lib/python3.12/site-packages/torch/utils/data/dataloader.py:674: in _next_data
    index = self._next_index()  # may raise StopIteration
.venv/lib/python3.12/site-packages/torch/utils/data/dataloader.py:621: in _next_index
    return next(self._sampler_iter)  # may raise StopIteration
.venv/lib/python3.12/site-packages/torch/utils/data/sampler.py:287: in __iter__
    for idx in self.sampler:
.venv/lib/python3.12/site-packages/torch/utils/data/sampler.py:167: in __iter__
    yield from torch.randperm(n, generator=generator).tolist()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <torch.utils._device.DeviceContext object at 0x734bb810f110>
func = <built-in method randperm of type object at 0x734c24f887e0>, types = ()
args = (100,)
kwargs = {'device': device(type='cuda'), 'generator': <torch._C.Generator object at 0x734b2cab4750>}

    def __torch_function__(self, func, types, args=(), kwargs=None):
        kwargs = kwargs or {}
        if func in _device_constructors() and kwargs.get('device') is None:
            kwargs['device'] = self.device
>       return func(*args, **kwargs)
E       RuntimeError: Expected a 'cuda' device type for generator but found 'cpu'

.venv/lib/python3.12/site-packages/torch/utils/_device.py:78: RuntimeError

[lululxvi]

The only thing DeepXDE modifies when loading is this https://deepxde.readthedocs.io/en/latest/user/installation.html#pytorch-backend . Maybe this is the reason? Does everything work well in CPU?

[JakobEliasWagner]

Yes everything works well in CPU. Only whenever I have a GPU available it fails the test. Yes I think the pytorch backend setting the default tensor type might cause this. Is there a way for me to disable this during testing?

[lululxvi]

There is no API to disable this. You could add the following code right after importing deepxde

torch.set_default_device("cpu")

[JakobEliasWagner]

Okay thank you for your answer.

I believe there is currently 3 different ways to solve this issue:

1. I personally isolate the deepxde dependency during testing.

def test_deepxde_specific():
    import deepxde as dde

2. Set the device context after importing DeepXDE (as suggested by @lululxvi)
3. Or use environment variables to control device allocation.