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