Which device is expected by a Module?

1

hi there,

for some reason, I need to know the device that is expected by a Module,
is there a way to do that ?

2

That’s a relative question. Mistakes happen when in an operation which makes use of two tensors. Thus, from the pov of tensor 1, expected device would be tensor2’s, whereas from the pov tensor 2 it would be the other way around.

Don’t think about “modules” because errors are thrown by tensors inside modules, not modules themselves.

The closest thing you can do is, if you know which operation implies your module/tensor, to query the analogous module/tensor by calling tensor.device() or module.device().

3

Yes, sure, but I have to think about modules in my case, I’m handling them as black boxes

So I guess the answer is no and I must rely on the user to provide me with a device parameter

4

You can use a try except scheme… but does not seem to be the best way.

5

For modules with parameters, next(m.parameters()).device. Modules without parameters should take tensors on any device.
This will fail for modules with mixed devices, but you probably don’t expect those to work, do you?

Best regards

Thomas

6

thanks,

that’s clear that I could indeed get the device of some arbitrary parameter.
I guess the question becomes: is there some introspection trick to get the first parameter that will be involved in the forward function of a module, to then bring my tensor to that place.

(likewise, having the expected input shape would be cool, and is a related question)

anyways, for now I’m getting user instructions, at he cost of more parameters