# Buggy code
def _register_forward_hooks(self):
    """Add hook before forward propagation of each layer to block forward computation until the push-pull and
    parameter update is finished. The blocking is implemented using a lock."""
    # Recursively find all submodules
    submodules = []
    q = queue.LifoQueue()
    for mod in self._model.children():
        q.put(mod)
    while not q.empty():
        mod = q.get()
        if len(list(mod.children())) == 0:
            submodules.append(mod)
        else:
            for m in mod.children():
                q.put(m)

    def pre_forward_hook(mod, input):
        for p in mod.parameters():
            if p in self._handles:
                del self._handles[p]
            if p not in self._locks:
                continue
            with self._locks[p]:
                self._logger.debug(
                    "{} {} is ready.".format(self._desc, self._parameter_names[p])
                )

        self._logger.debug("{} starts forward {}.".format(self._desc, mod))

    def after_forward_hook(mod, input, result):
        self._logger.debug("{} finished forward {}.".format(self._desc, mod))

    # Register pre-hook and hook for each module
    for mod in reversed(submodules):
        self._logger.debug(
            "{} registers forward hook on module {}".format(self._desc, mod)
        )
        mod.register_forward_pre_hook(pre_forward_hook)
        mod.register_forward_hook(after_forward_hook)


# 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 _register_forward_hooks(self):
    """Add hook before forward propagation of each layer to block forward computation until the push-pull and
    parameter update is finished. The blocking is implemented using a lock."""
    # Recursively find all submodules
    submodules = []
    q = queue.LifoQueue()
    for mod in self._model.children():
        q.put(mod)
    while not q.empty():
        mod = q.get()
        if len(list(mod.children())) == 0:
            submodules.append(mod)
        else:
            for m in mod.children():
                q.put(m)

    def pre_forward_hook(mod, input):
        for p in mod.parameters():
            if p in self._handles:
                del self._handles[p]
            if p not in self._locks:
                continue
            with self._locks[p]:
                self._logger.debug(
                    "{} {} is ready.".format(self._desc, self._parameter_names[id(p)])
                )

        self._logger.debug("{} starts forward {}.".format(self._desc, mod))

    def after_forward_hook(mod, input, result):
        self._logger.debug("{} finished forward {}.".format(self._desc, mod))

    # Register pre-hook and hook for each module
    for mod in reversed(submodules):
        self._logger.debug(
            "{} registers forward hook on module {}".format(self._desc, mod)
        )
        mod.register_forward_pre_hook(pre_forward_hook)
        mod.register_forward_hook(after_forward_hook)