# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/130_models.Hybrid.ipynb (unless otherwise specified).

__all__ = ['Hybrid']

# Cell
from ..imports import *
from .layers import *

# Cell
class Hybrid(Module):
    def __init__(self, dls, models, flatten=True, head=create_mlp_head, device=None, **kwargs):

        device = ifnone(device, default_device())
        self.models = nn.ModuleList()
        for m in models:
            m.head = Identity()
            self.models.append(m)
        self.flatten = Reshape(-1) if flatten else None
        with torch.no_grad():
            self.head = Noop
            self.head_nf = self.forward(first(dls.train)[0]).shape[-1]
        self.head = head(self.head_nf, dls.c, **kwargs)
        self.to(device=device)

    def forward(self, xs):
        out = []
        for x,m in zip(xs, self.models):
            _out = m(*x) if isinstance(x, L) else m(x)
            if self.flatten is not None and _out.ndim == 3: _out = self.flatten(_out)
            out = _out if out == [] else torch.cat([out, _out], dim=1)
        return self.head(out)