You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
TST check for device mismatch in solvers that support GPU input #852
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, cputorch.set_default_device("meta") # bare nx.ones(...) now defaults to metatry:
M=ot.batch.dist_batch(X, X)
res=ot.solve_batch(M, reg=0.1, max_iter=10, method="sinkhorn")
assertres.plan.device.type=="cpu"# fails/raises if any alloc leaks to metafinally:
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
[MRG] Check for device mismatch with a meta default device and fix ba…