我正在使用代码（笔记本）来预测。现在我应该清理代码并将其用于生产（更改为.py文件并导入等）。我正在为我的模型使用一个类，如下所示：

class Autoencoder(nn.Module):

    def __init__(self, input_dim, par_dim):
        super().__init__()

    def encode(self,y):
    return x

    def forward():....
    def test():...

培训批次和培训如下（它们不在类中）：

def train_batch(model, optimizer, device, batch, labels):

def train(model, device, epochs, train_iterator, optimizer, validate_iterator):

然后我考虑模型为：

model              = Autoencoder()

optimizer          = torch.optim.Adam(model.parameters() )

loss = train(model, device, epochs, train_iterator, optimizer, validate_iterator)

所以，我的问题是：我应该为类？使用一个.py文件，为train使用一个.py文件，为train_batch使用一个.py文件？
或者我可以在类内部使用这个train和train_batch？
你有任何教程（视频或stackoverflow链接）为这类工作？谢谢。谢谢