| 340 | return outputs |
| 341 | |
| 342 | class CostRegNet(nn.Module): |
| 343 | def __init__(self, in_channels, base_channels,stage=0): |
| 344 | super(CostRegNet, self).__init__() |
| 345 | self.cosR_small=CostRegNet_part(in_channels, base_channels,stage=0) |
| 346 | self.cosR_huge=CostRegNet_part(in_channels, base_channels,stage=0) |
| 347 | def forward(self, x): |
| 348 | results=torch.cat((self.cosR_small(x),self.cosR_huge(x)),axis=1) |
| 349 | return results |
| 350 | class CostRegNet_refine(nn.Module): |
| 351 | def __init__(self, in_channels, base_channels,stage=0): |
| 352 | super(CostRegNet_refine, self).__init__() |