CODEWITHSUNDEEP
unit-testingtestingpytorch
ikamen
ikamen

Reputation: 3485

How to test pytorch GPU code on a CPU machine

I want to verify on CPU machine that I've moved all tensors to CUDA device correctly. Is there a way to create a mock device that is still computed on CPU but marked to be incompatible with tensors I forget to move to this device?

Motivation: this would speed up my development cycle as I don't need to deploy code and run it on the GPU server only to find out I forgot to specify device on some tensor / model.

Upvotes: 2

Views: 772

Answers (1)

ikamen
ikamen

Reputation: 3485

Maybe the meta device is what you need.

>>> a = torch.ones(2) # defaults to cpu
>>> b = torch.ones(2, device='meta')
>>> c = a + b
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, meta and cpu!

Note that original intention for meta device is to compute shape and dtype of output without actually doing full computation. Tensors on meta device have no data in them, so I am not sure if there are some aspects of the test where meta device behaves different from a "real" device.

Interestingly I couldn't find much documentation about this meta device at all. I found it mentioned here https://github.com/pytorch/pytorch/issues/61654#issuecomment-879989145

Upvotes: 1

Related Questions