benchmarks/dlrm/ootb/dlrm_s_caffe2.py [834:857]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        if self.weighted_pooling == "learned":
            for i, w in enumerate(self.emb_vw):
                # select device
                if self.ndevices > 1:
                    d = i % self.ndevices
                # create tags
                on_device = "" if self.ndevices <= 1 else "gpu_" + str(d) + "/"
                _tag_one = on_device + tag_one
                _tag_lr = on_device + tag_lr
                # pickup gradient
                w_grad = self.gradientMap[w]
                # update weights
                if self.ndevices > 1:
                    with core.DeviceScope(
                        core.DeviceOption(workspace.GpuDeviceType, d)
                    ):
                        self.model.ScatterWeightedSum(
                            [w, _tag_one, w_grad.indices,
                            w_grad.values, _tag_lr], w
                        )
                else:
                    self.model.ScatterWeightedSum(
                        [w, _tag_one, w_grad.indices, w_grad.values, _tag_lr], w
                    )
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



benchmarks/dlrm/ootb/dlrm_s_caffe2.py [1003:1026]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        if self.weighted_pooling == "learned":
            for i, w in enumerate(self.emb_vw):
                # select device
                if self.ndevices > 1:
                    d = i % self.ndevices
                # create tags
                on_device = "" if self.ndevices <= 1 else "gpu_" + str(d) + "/"
                _tag_one = on_device + tag_one
                _tag_lr = on_device + tag_lr
                # pickup gradient
                w_grad = self.gradientMap[w]
                # update weights
                if self.ndevices > 1:
                    with core.DeviceScope(
                        core.DeviceOption(workspace.GpuDeviceType, d)
                    ):
                        self.model.ScatterWeightedSum(
                            [w, _tag_one, w_grad.indices,
                            w_grad.values, _tag_lr], w
                        )
                else:
                    self.model.ScatterWeightedSum(
                        [w, _tag_one, w_grad.indices, w_grad.values, _tag_lr], w
                    )
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -