# Buggy code
def _adam(self, p):
    """Performs a single optimization step using Adam optimizer on a parameter.
    Arguments:
        p: The parameter to be updated.
    """
    for group in self.param_groups:
        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
            grad = p.grad.data
            if grad.is_sparse:
                raise RuntimeError(
                    "Adam does not support sparse gradients, please consider SparseAdam instead"
                )
            amsgrad = group["amsgrad"]

            state = self.state[p]

            # State initialization
            if len(state) == 0:
                state["step"] = 0

                # Exponential moving average of gradient values
                state["exp_avg"] = torch.zeros_like(p.data)

                # Exponential moving average of squared gradient values
                state["exp_avg_sq"] = torch.zeros_like(p.data)
                if amsgrad:
                    # Maintains max of all exp. moving avg. of sq. grad. values
                    state["max_exp_avg_sq"] = torch.zeros_like(p.data)

            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
            if amsgrad:
                max_exp_avg_sq = state["max_exp_avg_sq"]
            beta1, beta2 = group["betas"]

            state["step"] += 1

            if group["weight_decay"] != 0:
                grad.add_(group["weight_decay"], p.data)

            # Decay the first and second moment running average coefficient
            exp_avg.mul_(beta1).add_(1 - beta1, grad)
            exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
            if amsgrad:
                # Maintains the maximum of all 2nd moment running avg. till now
                torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)

                # Use the max. for normalizing running avg. of gradient
                denom = max_exp_avg_sq.sqrt().add_(group["eps"])
            else:
                denom = exp_avg_sq.sqrt().add_(group["eps"])

            bias_correction1 = 1 - beta1 ** state["step"]
            bias_correction2 = 1 - beta2 ** state["step"]
            step_size = group["lr"] * math.sqrt(bias_correction2) / bias_correction1

            p.data.addcdiv_(-step_size, exp_avg, denom)
            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 _adam(self, p):
    """Performs a single optimization step using Adam optimizer on a parameter.
    Arguments:
        p: The parameter to be updated.
    """
    for group in self.param_groups:
        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
            grad = p.grad.data
            if grad.is_sparse:
                raise RuntimeError(
                    "Adam does not support sparse gradients, please consider SparseAdam instead"
                )
            amsgrad = group["amsgrad"]

            state = self.state[p]

            # State initialization
            if len(state) == 0:
                state["step"] = 0

                # Exponential moving average of gradient values
                state["exp_avg"] = torch.zeros_like(p.data)

                # Exponential moving average of squared gradient values
                state["exp_avg_sq"] = torch.zeros_like(p.data)
                if amsgrad:
                    # Maintains max of all exp. moving avg. of sq. grad. values
                    state["max_exp_avg_sq"] = torch.zeros_like(p.data)

            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
            if amsgrad:
                max_exp_avg_sq = state["max_exp_avg_sq"]
            beta1, beta2 = group["betas"]

            state["step"] += 1

            if group["weight_decay"] != 0:
                grad.add_(group["weight_decay"], p.data)

            # Decay the first and second moment running average coefficient
            exp_avg.mul_(beta1).add_(1 - beta1, grad)
            exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
            if amsgrad:
                # Maintains the maximum of all 2nd moment running avg. till now
                torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)

                # Use the max. for normalizing running avg. of gradient
                denom = max_exp_avg_sq.sqrt().add_(group["eps"])
            else:
                denom = exp_avg_sq.sqrt().add_(group["eps"])

            bias_correction1 = 1 - beta1 ** state["step"]
            bias_correction2 = 1 - beta2 ** state["step"]
            step_size = group["lr"] * math.sqrt(bias_correction2) / bias_correction1

            p.data.addcdiv_(-step_size, exp_avg, denom)
            break