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.