How can I add a function(def forward_per_GAP:) to the class(model_cls)?

[HYC986]

Now I am running a multi-trial NAS code. But I was confused when I using the model_cls.
Is it possible to add function(def forward_per_GAP:) to the class model_cls?
I just hope to can use a = model_cls.forward_per_GAP(x) in my code.
But i didn't make it.

@model_wrapper
class Net(nn.Module):
def init(cls):
super().init()

    conv1_choice1 = nn.ValueChoice([32, 64, 128, 256])
    conv1_choice2 = nn.ValueChoice([32, 64, 128, 256])
    cls.features = nn.Sequential(
        nn.Conv2d(1, conv1_choice1, 3, 1),
        nn.ReLU(),
        nn.Conv2d(conv1_choice1, conv1_choice2, 3, 1),
        nn.ReLU(),
        nn.LayerChoice([nn.Conv2d(conv1_choice2, 64, 3, 1), DepthwiseSeparableConv(conv1_choice2, 64)]),
        nn.MaxPool2d(2),
        nn.Dropout(0.5),
        nn.AdaptiveAvgPool2d((10, 10)),
        nn.Flatten(),
    )
    cls.classifier = nn.Sequential(
        nn.Linear(6400, 128),
        nn.ReLU(),
        nn.Dropout(0.5),
        nn.Linear(128, 10),
    )
def forward(cls, x):
    x = cls.features(x)
    x = cls.classifier(x)
    return x

def forward_pre_GAP(cls, x):                      **# i want to use this function in mode_cls. But don't know how to do.**
    output = x
    output = cls.features(output)
    return output

I am appreciate for any advice.