apps/model.py

import pytorch_lightning as pl
from apps.project_model2 import UNet

class Segmenter(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = UNet()
    def forward(self, data):
        pred = self.model(data)
        return pred
model=Segmenter()