Hy guys, I have two model.
The first is a compound model with fan_resnet50 and a reset ad hoc with 4 basicBlock.
The second model is a resnet50 pretrained.
I must remove the FC layer of these models and do a torch.cat of this 2 model.
> class Identity(nn.Module):
>     def __init__(self):
>         super().__init__()
> 
>     def forward(self, x):
>         return x
> class Compound_Model_final(nn.Module):
>     def __init__(self):
>         super().__init__()
>         self.model1 = fcn_resnet50(num_classes=9)
>         self.model2= _resnet_shallow(pretrained=False)
>         self.model3 = models.resnet50(pretrained=False)
>         self.model3.load_state_dict(torch.load("Regressione_ResNet50_cartesian_angle_M95.pth"))
>     def forward(self, x):
>         dist = self.model1(x)['out']
>         #dist = torch.softmax(dist.view(dist.shape[0],-1), dim =1).view(dist.shape)
>         n_features_x = self.model2.fc.in_features
>         n_features_y =self.model3.fc.in_features
>         self.model2.fc = Identity()
>         self.model3.fc = Identity()
> 
>         z = torch.cat([self.model2.fc, self.model3.fc ],1)
>         xyuv=nn.Linear((128*(n_features_x + n_features_y)), 4)
> 
> 
> 
> 
>         return dist, xyuv

i do this wrapper class. Is it right for you?