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 <class 'main.MyModel'>

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

f-dangel / backpack

Support for Custom models? #269