# 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.
    """

    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():
            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 p in nn_self.parameters():
            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 (RuntimeError)
import syft
import torch
from torchvision import models
import torch.nn as nn

hook = syft.TorchHook(torch)
worker = syft.VirtualWorker(hook, id="worker")

model = models.alexnet(pretrained=True)
for param in model.parameters():
param.requires_grad=False
model.classifier[6] = nn.Linear(model.classifier[6].in_features, 3)
model.send(worker)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-15-a250859d9a13> in <module>
----> 1 model.send(worker)

~/implementation/PyGrid/gateway/src/syft/syft/frameworks/torch/hook/hook.py in module_send_(nn_self, force_send, *dest, **kwargs)
608
609             if module_is_missing_grad(nn_self):
--> 610                 create_grad_objects(nn_self)
611
612             for p in nn_self.parameters():

~/implementation/PyGrid/gateway/src/syft/syft/frameworks/torch/hook/hook.py in create_grad_objects(model)
600             for p in model.parameters():
601                 o = p.sum()
--> 602                 o.backward()
603                 if p.grad is not None:
604                     p.grad -= p.grad

~/implementation/PyGrid/gateway/src/syft/syft/generic/frameworks/hook/trace.py in trace_wrapper(*args, **kwargs)
81                 syft.hook.trace.logs.append((command, response))
82             else:
---> 83                 response = func(*args, **kwargs)
84
85             return response

~/implementation/PyGrid/gateway/src/syft/syft/generic/frameworks/hook/hook.py in overloaded_native_method(self, *args, **kwargs)
436                 except BaseException as e:
437                     # we can make some errors more descriptive with this method
--> 438                     raise route_method_exception(e, self, args, kwargs)
439
440             else:  # means that there is a wrapper to remove

~/implementation/PyGrid/gateway/src/syft/syft/generic/frameworks/hook/hook.py in overloaded_native_method(self, *args, **kwargs)
432
433                 try:
--> 434                     response = method(*args, **kwargs)
435
436                 except BaseException as e:

~/anaconda3/lib/python3.7/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
193                 products. Defaults to ``False``.
194         """
--> 195         torch.autograd.backward(self, gradient, retain_graph, create_graph)
196
197     def register_hook(self, hook):

~/anaconda3/lib/python3.7/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
97     Variable._execution_engine.run_backward(
98         tensors, grad_tensors, retain_graph, create_graph,
---> 99         allow_unreachable=True)  # allow_unreachable flag
100
101

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

# 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.
    """

    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 p in nn_self.parameters():
            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