# Buggy code
def _hook_module(self):
    """Overloading torch.nn.Module with PySyft functionality, the primary module
    responsible for core ML functionality such as Neural network layers and
    loss functions.
    It is important to note that all the operations are actually in-place.
    """
    self.element_iter_dict = {}

    def register_element_iterator(name, func):
        """register an internal element buffer iterator"""
        if name in self.element_iter_dict.keys():
            return
        self.element_iter_dict[name] = func

    def tensor_iterator(nn_self):
        """adding relavant iterators for the tensor elements"""
        iterators = [
            "parameters",
            "buffers",
        ]  # all the element iterators from nn module should be listed here,
        return [getattr(nn_self, iter) for iter in iterators]

    def module_is_missing_grad(model):
        """Checks if all the parameters in the model have been assigned a gradient"""
        for p in model.parameters():
            if p.grad is None:
                return True
        return False

    def create_grad_objects(model):
        """Assigns gradient to model parameters if not assigned"""
        for p in model.parameters():
            if p.requires_grad:  # check if the object requires a grad object
                o = p.sum()
                o.backward()
                if p.grad is not None:
                    p.grad -= p.grad

    def module_send_(nn_self, *dest, force_send=False, **kwargs):
        """Overloads torch.nn instances so that they could be sent to other workers"""

        if module_is_missing_grad(nn_self):
            create_grad_objects(nn_self)

        for element_iter in tensor_iterator(nn_self):
            for p in element_iter():
                p.send_(*dest, **kwargs)

        if isinstance(nn_self.forward, Plan):
            nn_self.forward.send(*dest, force=force_send)

        return nn_self

    self.torch.nn.Module.send = module_send_
    self.torch.nn.Module.send_ = module_send_

    def module_move_(nn_self, destination):
        params = list(nn_self.parameters())
        for p in params:
            p.move_(destination)

    self.torch.nn.Module.move = module_move_

    # def module_end_get_(nn_self):
    #     """Overloads send to remote for torch.nn.Module."""
    #     if module_is_missing_grad(nn_self):
    #         create_grad_objects(nn_self)
    #
    #     for p in nn_self.parameters():
    #         p.end_get()
    #
    #     return nn_self
    #
    # self.torch.nn.Module.end_get = module_end_get_
    #
    # def module_move_(nn_self, dest):
    #     return nn_self.send(dest).end_get()
    #
    # self.torch.nn.Module.move = module_move_

    def module_get_(nn_self):
        """overloads torch.nn instances with get method so that parameters could be sent back to owner"""
        for p in nn_self.parameters():
            p.get_()

        if isinstance(nn_self.forward, Plan):
            nn_self.forward.get()

        return nn_self

    self.torch.nn.Module.get_ = module_get_
    self.torch.nn.Module.get = module_get_

    def module_share_(nn_self, *args, **kwargs):
        """Overloads fix_precision for torch.nn.Module."""
        # TODO: add .data and .grad to syft tensors
        if module_is_missing_grad(nn_self):
            create_grad_objects(nn_self)

        for p in nn_self.parameters():
            p.share_(*args, **kwargs)

        return nn_self

    self.torch.nn.Module.share_ = module_share_
    self.torch.nn.Module.share = module_share_

    def module_fix_precision_(nn_self, *args, **kwargs):
        """Overloads fix_precision for torch.nn.Module."""
        if module_is_missing_grad(nn_self):
            create_grad_objects(nn_self)

        for p in nn_self.parameters():
            p.fix_precision_(*args, **kwargs)

        return nn_self

    self.torch.nn.Module.fix_precision_ = module_fix_precision_
    self.torch.nn.Module.fix_precision = module_fix_precision_
    self.torch.nn.Module.fix_prec = module_fix_precision_

    def module_float_precision_(nn_self):
        """Overloads float_precision for torch.nn.Module, convert fix_precision
        parameters to normal float parameters"""
        # TODO: add .data and .grad to syft tensors
        # if module_is_missing_grad(nn_self):
        #    create_grad_objects(nn_self)

        for p in nn_self.parameters():
            p.float_precision_()

        return nn_self

    self.torch.nn.Module.float_precision_ = module_float_precision_
    self.torch.nn.Module.float_precision = module_float_precision_
    self.torch.nn.Module.float_prec = module_float_precision_

    def module_copy(nn_self):
        """Returns a copy of a torch.nn.Module"""
        return copy.deepcopy(nn_self)

    self.torch.nn.Module.copy = module_copy

    @property
    def owner(nn_self):
        for p in nn_self.parameters():
            return p.owner

    self.torch.nn.Module.owner = owner

    @property
    def location(nn_self):
        try:
            for p in nn_self.parameters():
                return p.location
        except AttributeError:
            raise AttributeError(
                "Module has no attribute location, did you already send it to some location?"
            )

    self.torch.nn.Module.location = location

    # Make sure PySyft uses the PyTorch version
    self.torch.nn.modules.rnn._rnn_impls["LSTM"] = self.torch.lstm

    # Add support for GRUs
    self.torch.nn.modules.rnn._rnn_impls["GRU"] = self.torch.gru

    # Override _VF.LSTM_Cell and _VF.GRU_Cell with torch.LSTM_Cell and torch.GRU_Cell
    # With the pytorch-based version
    self.torch.nn.modules.rnn._VF = self.torch


# Traceback (PureTorchTensorFoundError)
WARNING: Logging before flag parsing goes to stderr.
W0813 08:46:34.190028 140520788572032 secure_random.py:26] Falling back to insecure randomness since the required custom op could not be found for the installed version of TensorFlow. Fix this by compiling custom ops. Missing file was '/usr/local/lib/python3.6/dist-packages/tf_encrypted/operations/secure_random/secure_random_module_tf_1.14.0.so'
W0813 08:46:34.211576 140520788572032 deprecation_wrapper.py:119] From /usr/local/lib/python3.6/dist-packages/tf_encrypted/session.py:26: The name tf.Session is deprecated. Please use tf.compat.v1.Session instead.

0it [00:00, ?it/s]

Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ../data/cifar-10-python.tar.gz

170500096it [00:04, 40609204.86it/s]

---------------------------------------------------------------------------

PureTorchTensorFoundError                 Traceback (most recent call last)

/content/PySyft/syft/frameworks/torch/tensors/interpreters/native.py in handle_func_command(cls, command)
300             new_args, new_kwargs, new_type, args_type = syft.frameworks.torch.hook_args.unwrap_args_from_function(
--> 301                 cmd, args, kwargs, return_args_type=True
302             )

16 frames

PureTorchTensorFoundError:


During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)

/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in batch_norm(input, running_mean, running_var, weight, bias, training, momentum, eps)
1695     return torch.batch_norm(
1696         input, weight, bias, running_mean, running_var,
-> 1697         training, momentum, eps, torch.backends.cudnn.enabled
1698     )
1699

RuntimeError: running_mean should contain 3 elements not 32

# Fixed code
def _hook_module(self):
    """Overloading torch.nn.Module with PySyft functionality, the primary module
    responsible for core ML functionality such as Neural network layers and
    loss functions.
    It is important to note that all the operations are actually in-place.
    """
    self.element_iter_dict = {}

    def register_element_iterator(name, func):
        """register an internal element buffer iterator"""
        if name in self.element_iter_dict.keys():
            return
        self.element_iter_dict[name] = func

    def tensor_iterator(nn_self):
        """adding relavant iterators for the tensor elements"""
        iterators = [
            "parameters",
            "buffers",
        ]  # all the element iterators from nn module should be listed here,
        return [getattr(nn_self, iter) for iter in iterators]

    def module_is_missing_grad(model):
        """Checks if all the parameters in the model have been assigned a gradient"""
        for p in model.parameters():
            if p.grad is None:
                return True
        return False

    def create_grad_objects(model):
        """Assigns gradient to model parameters if not assigned"""
        for p in model.parameters():
            if p.requires_grad:  # check if the object requires a grad object
                o = p.sum()
                o.backward()
                if p.grad is not None:
                    p.grad -= p.grad

    def module_send_(nn_self, *dest, force_send=False, **kwargs):
        """Overloads torch.nn instances so that they could be sent to other workers"""

        if module_is_missing_grad(nn_self):
            create_grad_objects(nn_self)

        for element_iter in tensor_iterator(nn_self):
            for p in element_iter():
                p.send_(*dest, **kwargs)

        if isinstance(nn_self.forward, Plan):
            nn_self.forward.send(*dest, force=force_send)

        return nn_self

    self.torch.nn.Module.send = module_send_
    self.torch.nn.Module.send_ = module_send_

    def module_move_(nn_self, destination):
        params = list(nn_self.parameters())
        for p in params:
            p.move_(destination)

    self.torch.nn.Module.move = module_move_

    # def module_end_get_(nn_self):
    #     """Overloads send to remote for torch.nn.Module."""
    #     if module_is_missing_grad(nn_self):
    #         create_grad_objects(nn_self)
    #
    #     for p in nn_self.parameters():
    #         p.end_get()
    #
    #     return nn_self
    #
    # self.torch.nn.Module.end_get = module_end_get_
    #
    # def module_move_(nn_self, dest):
    #     return nn_self.send(dest).end_get()
    #
    # self.torch.nn.Module.move = module_move_

    def module_get_(nn_self):
        """overloads torch.nn instances with get method so that parameters could be sent back to owner"""
        for element_iter in tensor_iterator(nn_self):
            for p in element_iter():
                p.get_()

        if isinstance(nn_self.forward, Plan):
            nn_self.forward.get()

        return nn_self

    self.torch.nn.Module.get_ = module_get_
    self.torch.nn.Module.get = module_get_

    def module_share_(nn_self, *args, **kwargs):
        """Overloads fix_precision for torch.nn.Module."""
        if module_is_missing_grad(nn_self):
            create_grad_objects(nn_self)

        for element_iter in tensor_iterator(nn_self):
            for p in element_iter():
                p.share_(*args, **kwargs)

        return nn_self

    self.torch.nn.Module.share_ = module_share_
    self.torch.nn.Module.share = module_share_

    def module_fix_precision_(nn_self, *args, **kwargs):
        """Overloads fix_precision for torch.nn.Module."""
        if module_is_missing_grad(nn_self):
            create_grad_objects(nn_self)

        for element_iter in tensor_iterator(nn_self):
            for p in element_iter():
                p.fix_precision_(*args, **kwargs)

        return nn_self

    self.torch.nn.Module.fix_precision_ = module_fix_precision_
    self.torch.nn.Module.fix_precision = module_fix_precision_
    self.torch.nn.Module.fix_prec = module_fix_precision_

    def module_float_precision_(nn_self):
        """Overloads float_precision for torch.nn.Module, convert fix_precision
        parameters to normal float parameters"""
        # TODO: add .data and .grad to syft tensors
        # if module_is_missing_grad(nn_self):
        #    create_grad_objects(nn_self)

        for element_iter in tensor_iterator(nn_self):
            for p in element_iter():
                p.float_precision_()

        return nn_self

    self.torch.nn.Module.float_precision_ = module_float_precision_
    self.torch.nn.Module.float_precision = module_float_precision_
    self.torch.nn.Module.float_prec = module_float_precision_

    def module_copy(nn_self):
        """Returns a copy of a torch.nn.Module"""
        return copy.deepcopy(nn_self)

    self.torch.nn.Module.copy = module_copy

    @property
    def owner(nn_self):
        for p in nn_self.parameters():
            return p.owner

    self.torch.nn.Module.owner = owner

    @property
    def location(nn_self):
        try:
            for p in nn_self.parameters():
                return p.location
        except AttributeError:
            raise AttributeError(
                "Module has no attribute location, did you already send it to some location?"
            )

    self.torch.nn.Module.location = location

    # Make sure PySyft uses the PyTorch version
    self.torch.nn.modules.rnn._rnn_impls["LSTM"] = self.torch.lstm

    # Add support for GRUs
    self.torch.nn.modules.rnn._rnn_impls["GRU"] = self.torch.gru

    # Override _VF.LSTM_Cell and _VF.GRU_Cell with torch.LSTM_Cell and torch.GRU_Cell
    # With the pytorch-based version
    self.torch.nn.modules.rnn._VF = self.torch