from .unet import UNet
from .BottomHalf import BottomHalf
from .segnet import SegNet
import torch.nn as nn

class KeyModel(nn.Module):
    """A dummy model that simply returns a key from its dictionary input"""
    def __init__(self, pred_y_key: str):
        self.name = pred_y_key
        super().__init__()
    def forward(self, batch):
        return batch[self.name]


from .unet import UNet
from .BottomHalf import BottomHalf
from .segnet import SegNet
import torch.nn as nn

class KeyModel(nn.Module):
    """A dummy model that simply returns a key from its dictionary input"""
    def __init__(self, pred_y_key: str):
        self.name = pred_y_key
        super().__init__()
    def forward(self, batch):
        return batch[self.name]


A dummy model that simply returns a key from its dictionary input