Hello,
I want to change the model layers of my previously trained models. For example, I have a MobileNet trained on my data, and now I want to remove the FC layers. For this purpose, I replace the FC layer with a class called Identity, which you can find the definition below. (I took this suggestion from one of the admins in here, but I do not remember the post)
class Identity(nn.Module):
    def __init__(self):
        super(Identity, self).__init__()

    def forward(self, x):
        return x

The procedures are this:

Define a MobileNet by using torchvision.model

Move the model to nn.DataParallel

Load the previously trained model
Replace FC layers with Identity class

Now, the problem here is that the model is not fully replaced, and I think it is because there are other copies on other GPUs as well.
Is there any solution to change all the models? This is the result of print(model).

DataParallel(
(module): MobileNetV2(
(features): Sequential(
…
…
)
)
(classifier): Sequential(
(0): Linear(in_features=1280, out_features=512, bias=True)
(1): ReLU()
(2): Dropout(p=0.2, inplace=False)
(3): Linear(in_features=512, out_features=7, bias=True)
)
)
(classifier): Identity()
)

You can see that there are two classifiers (FC layers) because I use 2 GPUs.
One More Important thing:
When I set the number of GPUs to 1, so I do not have the above problem, I get this error:
OpenBLAS Warning : Detect OpenMP Loop and this application may hang. Please rebuild the library with USE_OPENMP=1 option.