这是我的模型

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28 * 28, 512, dtype=torch.float),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

我正在跟踪Crate Pipeline的这个repo：https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/pytorch_gcs_data_training.ipynb
我目前正在做一个项目，涉及使用“torch-model-archiver”工具部署PyTorch模型。虽然我理解了使用“torch-model-archiver”创建.mar文件的基本过程，但我在使用该工具提供的默认处理程序时遇到了一些挑战。
我不太确定如何实现这一点。有人能指导我如何为“torch-model-archiver”创建一个自定义处理程序吗？我想使用这个自定义处理程序为我的PyTorch模型生成一个.mar文件。