MCPcopy Create free account
hub / github.com/pytorch/pytorch / _get_param_to_device

Function _get_param_to_device

caffe2/python/optimizer.py:2039–2044  ·  view source on GitHub ↗
(model)

Source from the content-addressed store, hash-verified

2037
2038
2039def _get_param_to_device(model):
2040 # Infer blob devices by going through the net and param_init_net
2041 # ops and observing the device used to create or use the blob.
2042 param_to_device = core.InferBlobDevices(model.net)
2043 param_to_device.update(core.InferBlobDevices(model.param_init_net))
2044 return param_to_device
2045
2046
2047def get_param_device(param_name, grad, param_to_device=None, default_device=None):

Callers 1

_buildFunction · 0.85

Calls 1

updateMethod · 0.45

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…