How to use Pytorch to create a custom EfficientNet with the last layer written correctly
I have a classification problem to predict 8 classes for example, I am using EfficientNetB3 in pytorch from here. However, I got confused on whether my custom class is correctly written. I think I want to strip the last layer of the pre-trained model to suit the 8 outputs right? Did I do it correctly? Because when I print y_preds = model(images) in my DataLoader, it seems to give me 1536 predictions. Is this an expected behavior?
!pip install geffnet 
import geffnet

class EfficientNet(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.model = geffnet.create_model(config.effnet, pretrained=True)
        n_features = self.model.classifier.in_features
        # does the name fc matter?
        self.fc = nn.Linear(n_features, config.num_classes)
        self.model.classifier = nn.Identity()
        
    def extract(self, x):
        x = self.model(x)
        return x

    def forward(self, x):
        x = self.extract(x).squeeze(-1).squeeze(-1)
        return x
    
model = EfficientNet(config=config)
if torch.cuda.is_available():
    model.cuda()

Sample code for printing y_pred:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
for step, (images, labels) in enumerate(sample_loader):
    images = images.to(device)
    labels = labels.to(device)
    batch_size = images.shape[0]        
    y_preds = model(images)
    print('The predictions of the 4 images is as follows\n', y_preds)
    break


You're not even using self.fc in forward pass.
Either just introduce it as:
def forward(self, x):
    ....
    x = extract(x)...
    x = fc(x)
    return x

Or you can simply replace the layer named classifier (this way you don't need Identity layer):
self.model.classifier = nn.Linear(n_features, config.num_classes)

Also, here config.num_classes should be 8.
