# Buggy code
def _sgd(self, p):
    """Performs a single optimization step using SGD optimizer on a parameter.
    Arguments:
        p: The parameter to be updated.
    """
    for group in self.param_groups:
        weight_decay = group["weight_decay"]
        momentum = group["momentum"]
        dampening = group["dampening"]
        nesterov = group["nesterov"]

        for gp in group["params"]:
            if (
                self._parameter_names[p] != self._parameter_names[gp]
                or gp.shape != p.shape
            ):
                continue
            self._logger.debug(
                "{} is updating {}".format(self._desc, self._parameter_names[p])
            )
            if p.grad is None:
                continue
            d_p = p.grad.data
            if weight_decay != 0:
                d_p.add_(weight_decay, p.data)
            if momentum != 0:
                param_state = self.state[p]
                if "momentum_buffer" not in param_state:
                    buf = param_state["momentum_buffer"] = torch.zeros_like(p.data)
                    buf.mul_(momentum).add_(d_p)
                else:
                    buf = param_state["momentum_buffer"]
                    buf.mul_(momentum).add_(1 - dampening, d_p)
                if nesterov:
                    d_p = d_p.add(momentum, buf)
                else:
                    d_p = buf
            p.data.add_(-group["lr"], d_p)
            break


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/usr/local/byteps/example/pytorch/benchmark_bytescheduler.py", line 134, in <module>
timeit.timeit(benchmark_step, number=args.num_warmup_batches)
File "/opt/anaconda/lib/python3.7/timeit.py", line 232, in timeit
return Timer(stmt, setup, timer, globals).timeit(number)
File "/opt/anaconda/lib/python3.7/timeit.py", line 176, in timeit
timing = self.inner(it, self.timer)
File "<timeit-src>", line 6, in inner
File "/usr/local/byteps/example/pytorch/benchmark_bytescheduler.py", line 115, in benchmark_step
output = model(data)
File "/opt/anaconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 541, in __call__
result = self.forward(*input, **kwargs)
File "/opt/anaconda/lib/python3.7/site-packages/torchvision/models/resnet.py", line 196, in forward
x = self.conv1(x)
File "/opt/anaconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 533, in __call__
result = hook(self, input)
File "/opt/anaconda/lib/python3.7/site-packages/byteps-0.1.0-py3.7-linux-x86_64.egg/byteps/bytescheduler/torch/optimizer.py", line 206, in pre_forward_hook
self._logger.debug("{} {} is ready.".format(self._desc, self._parameter_names[p]))
RuntimeError: bool value of Tensor with more than one value is ambiguous

# Fixed code
def _sgd(self, p):
    """Performs a single optimization step using SGD optimizer on a parameter.
    Arguments:
        p: The parameter to be updated.
    """
    for group in self.param_groups:
        weight_decay = group["weight_decay"]
        momentum = group["momentum"]
        dampening = group["dampening"]
        nesterov = group["nesterov"]

        for gp in group["params"]:
            if (
                self._parameter_names[id(p)] != self._parameter_names[id(gp)]
                or gp.shape != p.shape
            ):
                continue
            self._logger.debug(
                "{} is updating {}".format(self._desc, self._parameter_names[id(p)])
            )
            if p.grad is None:
                continue
            d_p = p.grad.data
            if weight_decay != 0:
                d_p.add_(weight_decay, p.data)
            if momentum != 0:
                param_state = self.state[p]
                if "momentum_buffer" not in param_state:
                    buf = param_state["momentum_buffer"] = torch.zeros_like(p.data)
                    buf.mul_(momentum).add_(d_p)
                else:
                    buf = param_state["momentum_buffer"]
                    buf.mul_(momentum).add_(1 - dampening, d_p)
                if nesterov:
                    d_p = d_p.add(momentum, buf)
                else:
                    d_p = buf
            p.data.add_(-group["lr"], d_p)
            break