Spaces:
Running
on
T4
Running
on
T4
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)) | |