Is there a way to retrieve the device type and device ID of a model, from inside a custom JIT graph pass? In other words, if a customer does a model.to("cuda:0"), is there any way to retrieve “cuda, 0” from inside a graph pass? Please let me know.
This depends on the state of the graph.
For example, in a freshly traced graph, you have the values (mostly) with as complete tensor type and that includes the device. Similarly dimensioned tensors (with scalar type, dimension (1d 2d 3d etc) and device, requires grad) have that information, but more incomplete values don’t.
Best regards
Thomas
Can be done in a custom graph pass during inference. Graph pass only has reference to std::shared_ptr<Graph> . But assuming that this is during inference, is it possible to retrieve the device-type and device-id (in case of cuda)? This could be scripted model or traced model.