from typing import Iterable
import torch
from torch import nn


def module_converter(module: nn.Module, keys: Iterable[str]) ->nn.ModuleDict:
    base = nn.ModuleDict()
    for n, m in module.named_modules():
        if n in keys:
            base.update({n: m})
    return base


class Flatten(nn.Module):

    def __init__(self):
        super(Flatten, self).__init__()

    def forward(self, input: torch.Tensor):
        return input.reshape(input.size(0), -1)

40c396677caf431e5e0fc1845e11d7c90d2ef3bf
<mask token>


class Flatten(nn.Module):

    def __init__(self):
        super(Flatten, self).__init__()
    <mask token>

<mask token>


class Flatten(nn.Module):

    def __init__(self):
        super(Flatten, self).__init__()

    def forward(self, input: torch.Tensor):
        return input.reshape(input.size(0), -1)

<mask token>


def module_converter(module: nn.Module, keys: Iterable[str]) ->nn.ModuleDict:
    base = nn.ModuleDict()
    for n, m in module.named_modules():
        if n in keys:
            base.update({n: m})
    return base


class Flatten(nn.Module):

    def __init__(self):
        super(Flatten, self).__init__()

    def forward(self, input: torch.Tensor):
        return input.reshape(input.size(0), -1)

from typing import Iterable
import torch
from torch import nn


def module_converter(module: nn.Module, keys: Iterable[str]) ->nn.ModuleDict:
    base = nn.ModuleDict()
    for n, m in module.named_modules():
        if n in keys:
            base.update({n: m})
    return base


class Flatten(nn.Module):

    def __init__(self):
        super(Flatten, self).__init__()

    def forward(self, input: torch.Tensor):
        return input.reshape(input.size(0), -1)
