# Buggy code
def _rmsprop(self, p):
    """Performs a single optimization step using RMSprop 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("RMSprop does not support sparse gradients")
            state = self.state[p]

            # State initialization
            if len(state) == 0:
                state["step"] = 0
                state["square_avg"] = torch.zeros_like(p.data)
                if group["momentum"] > 0:
                    state["momentum_buffer"] = torch.zeros_like(p.data)
                if group["centered"]:
                    state["grad_avg"] = torch.zeros_like(p.data)

            square_avg = state["square_avg"]
            alpha = group["alpha"]

            state["step"] += 1

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

            square_avg.mul_(alpha).addcmul_(1 - alpha, grad, grad)

            if group["centered"]:
                grad_avg = state["grad_avg"]
                grad_avg.mul_(alpha).add_(1 - alpha, grad)
                avg = (
                    square_avg.addcmul(-1, grad_avg, grad_avg).sqrt().add_(group["eps"])
                )
            else:
                avg = square_avg.sqrt().add_(group["eps"])

            if group["momentum"] > 0:
                buf = state["momentum_buffer"]
                buf.mul_(group["momentum"]).addcdiv_(grad, avg)
                p.data.add_(-group["lr"], buf)
            else:
                p.data.addcdiv_(-group["lr"], grad, avg)
            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 _rmsprop(self, p):
    """Performs a single optimization step using RMSprop 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("RMSprop does not support sparse gradients")
            state = self.state[p]

            # State initialization
            if len(state) == 0:
                state["step"] = 0
                state["square_avg"] = torch.zeros_like(p.data)
                if group["momentum"] > 0:
                    state["momentum_buffer"] = torch.zeros_like(p.data)
                if group["centered"]:
                    state["grad_avg"] = torch.zeros_like(p.data)

            square_avg = state["square_avg"]
            alpha = group["alpha"]

            state["step"] += 1

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

            square_avg.mul_(alpha).addcmul_(1 - alpha, grad, grad)

            if group["centered"]:
                grad_avg = state["grad_avg"]
                grad_avg.mul_(alpha).add_(1 - alpha, grad)
                avg = (
                    square_avg.addcmul(-1, grad_avg, grad_avg).sqrt().add_(group["eps"])
                )
            else:
                avg = square_avg.sqrt().add_(group["eps"])

            if group["momentum"] > 0:
                buf = state["momentum_buffer"]
                buf.mul_(group["momentum"]).addcdiv_(grad, avg)
                p.data.add_(-group["lr"], buf)
            else:
                p.data.addcdiv_(-group["lr"], grad, avg)
            break