model_arch.py

import torch
import segmentation_models_pytorch as smp

def get_model(weights_path=None):
    model = smp.Unet(
        encoder_name="resnet34",
        in_channels=11,
        classes=1
    )
    if weights_path:
        model.load_state_dict(torch.load(weights_path, map_location='cpu'))
    model.eval()
    return model