Function _ForEachDevice

caffe2/python/data_parallel_model.py:855–864

(devices, f, device_type, device_prefix, scoped=False,
                   *args, **kwargs)

def _ForEachDevice(devices, f, device_type, device_prefix, scoped=False,
                   *args, **kwargs):
    for device in devices:
        device_opt = core.DeviceOption(device_type, device)
        with core.DeviceScope(device_opt):
            if scoped:
                with core.NameScope("{}_{}".format(device_prefix, device)):
                    f(device, *args, **kwargs)
            else:
                f(device, *args, **kwargs)


def _AddGradientOperators(devices, model, losses_by_gpu):