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.
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
metadevice.Naively using
metaas 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:
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.