import torch
import torch.nn as nn
from torch import Tensor


class DeployC2f(nn.Module):

    def __init__(self, *args, **kwargs):
        super().__init__()

    def forward(self, x: Tensor) -> Tensor:
        x_main = self.main_conv(x)
        x_main = [x_main, x_main[:, self.mid_channels:, ...]]
        x_main.extend(blocks(x_main[-1]) for blocks in self.blocks)
        x_main.pop(1)
        return self.final_conv(torch.cat(x_main, 1))