How to create a torch.nn.Module model Class file?

I have a detecton2 model. For inference API we were using flask but now switching to Kfserving. Kfserving required a saved model file and the class name to load the model.

Previously I use this code snippet to load the model.

cfg = get_cfg()
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.8  # set threshold for this model
from detectron2.modeling import build_model
from detectron2.checkpoint import DetectionCheckpointer
model = build_model(cfg)

but Now I require something like this.

class TwoLayerNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
    def forward(self, x):

Were you able to figure it out?

KFServing (now named as KServe), has moved it’s support for PyTorch from the custom based model server to TorchServe. You can find more information about deploying a PyTorch model archive in Kserve here, and information about storing your model in MAR format here.