Skip to content

TST check for device mismatch in solvers that support GPU input #852

Description

@tomMoral

A common regression in various function is that they crash when input is on GPU due to internal buffer being allocated on CPU, leading to RuntimeError: Expected all tensors to be on the same device.
To systematically test function against this regression without GPU, one could use the meta device.
Naively using meta as input device does not work but a possible strategy is to keep real data on CPU and flip the default device to meta.
Any stray default-device allocation then lands on meta, mismatches the CPU input, and raises the device error — while the algorithm itself runs on real CPU data and should still work:

X = torch.randn(2, 4, 2)            # real data, cpu
torch.set_default_device("meta")   # bare nx.ones(...) now defaults to meta
try:
    M = ot.batch.dist_batch(X, X)
    res = ot.solve_batch(M, reg=0.1, max_iter=10, method="sinkhorn")
    assert res.plan.device.type == "cpu"   # fails/raises if any alloc leaks to meta
finally:
    torch.set_default_device("cpu")

We could audit the full code basis to run this on all functions that accept GPUs.
adding an helper fixture setting the default device to meta would make the test easy to setup/teardown.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions