Support for Custom models?

Hi,
I cannot find it in the docs for the package, but how is extend(model) actually implemented such that it computes the higher order extensions for the parameters of the model?
Basically if I want to calculate the Generalised Gauss Newton matrix using backpack for a model which doesn't actually use the torch.nn layers, but rather a custom Forward operation... is it possible ?

For example -->

```
class MyModel(nn.Module):
"""
Define the Forward Model of the Experiment
Args
-------
X (Tensor), Y (Tensor)

Returns
-------
Z : Tensor

"""
def __init__(self, X, Y ):

super().__init__()
self.X = torch.nn.Parameter(X)
self.Y = torch.nn.Parameter(Y)

def forward( self, X_0):
#some operation on X and Y
return ( self.X- X_0 )*self.Y
```

and then -->

```
x = torch.tensor([1.,2.,3.])
y = torch.tensor([11.,22.,33.])

x0 = torch.tensor([0.5])

inputs = torch.tensor([10.,18.,13.])

model = MyModel(x,y)
model = extend(model)
cost_function = extend(torch.nn.MSEloss())

preds = model(x0)
cost = cost_function(preds, inputs)
with backpack(extensions.GGNMP()):
cost.backward()

```

It doesn't work, and I get the following error -->

> Extension saving to ggnmp does not have an extension for Module

Should I modify the `extend(model)` somehow?
Another question, does backpack support Complex Tensors for computation of these higher order extensions?

Start by reproducing the provided MyModel example with extend(model), MSELoss, and extensions.GGNMP(), then read the implementation and documentation for custom Module support. Determine whether custom forward operations and complex tensors are supported; done should include a clear documented answer or a defined implementation and test for the requested behavior.

Written by the indexing model from the issue text.