# Buggy code
def get(self):
    """Just a pass through. This is most commonly used when calling .get() on a
    AutogradTensor which has also been shared."""
    self.child = self.child.get()
    # Remove the autograd node if a simple tensor is received
    if isinstance(self.child, torch.Tensor) and not self.child.is_wrapper:
        return self.child
    return self


# Traceback (AttributeError)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-1-db6dbbeffaa2> in <module>()
100     # Breaks
101     federated()
--> 102     encrypted() # breaks here - something about loss2.backward() causes the federated() demo to break
103 else:
104     # Works fine

<ipython-input-1-db6dbbeffaa2> in encrypted()
84
85     # 4) figure out which weights caused us to miss
---> 86     loss2.backward()
87
88 #     # 5) change those weights

/Users/atrask/anaconda/lib/python3.6/site-packages/syft-0.1.22a1-py3.6.egg/syft/frameworks/torch/hook/hook.py in overloaded_native_method(self, *args, **kwargs)
683                 # Put back the wrappers where needed
684                 response = syft.frameworks.torch.hook_args.hook_response(
--> 685                     method_name, response, wrap_type=type(self), new_self=self
686                 )
687

/Users/atrask/anaconda/lib/python3.6/site-packages/syft-0.1.22a1-py3.6.egg/syft/frameworks/torch/hook/hook_args.py in hook_response(attr, response, wrap_type, wrap_args, new_self)
243         response_hook_function = hook_method_response_functions[attr_id]
244         # Try running it
--> 245         new_response = response_hook_function(response)
246
247     except (IndexError, KeyError, AssertionError):  # Update the function in cas of an error

/Users/atrask/anaconda/lib/python3.6/site-packages/syft-0.1.22a1-py3.6.egg/syft/frameworks/torch/hook/hook_args.py in <lambda>(x)
502         f = many_fold
503
--> 504     return lambda x: f(lambdas, x)
505
506

/Users/atrask/anaconda/lib/python3.6/site-packages/syft-0.1.22a1-py3.6.egg/syft/frameworks/torch/hook/hook_args.py in two_fold(lambdas, args, **kwargs)
520
521 def two_fold(lambdas, args, **kwargs):
--> 522     return lambdas[0](args[0], **kwargs), lambdas[1](args[1], **kwargs)
523
524

/Users/atrask/anaconda/lib/python3.6/site-packages/syft-0.1.22a1-py3.6.egg/syft/frameworks/torch/hook/hook_args.py in <lambda>(i)
480         if isinstance(r, (list, tuple))  # if the rule is a list or tuple.
481         # Last if not, rule is probably == 1 so use type to return the right transformation.
--> 482         else lambda i: backward_func[wrap_type](i, **wrap_args)
483         for a, r in zip(response, rules)  # And do this for all the responses / rules provided
484     ]

/Users/atrask/anaconda/lib/python3.6/site-packages/syft-0.1.22a1-py3.6.egg/syft/frameworks/torch/hook/hook_args.py in <lambda>(i)
73 backward_func = {
74     TorchTensor: lambda i: i.wrap(),
---> 75     torch.Tensor: lambda i: i.wrap(),
76     torch.nn.Parameter: lambda i: torch.nn.Parameter(data=i),
77     PointerTensor: lambda i: i,

AttributeError: 'NoneType' object has no attribute 'wrap'

# Fixed code
def get(self):
    """Just a pass through. This is most commonly used when calling .get() on a
    AutogradTensor which has also been shared."""
    tensor = self.child.get()

    if isinstance(tensor, torch.Tensor):
        # Remove the autograd node if a simple tensor is received
        if not tensor.is_wrapper:
            return tensor
        # If it's a wrapper, then insert the autograd under the wrapper
        else:
            self.child = tensor.child
            tensor.child = self
            return tensor

    self.child = tensor
    return self