I changed the code a little bit and it started working . But the loss is still to high ; What should I make changes in model class to improve the accuracy ? I am Training 51000 images with a batch size of 256 sized at 224 x 224 , for 25 epochs.
Notebook Link
import torch.nn as nn
import torch.nn.functional as F

class EfficientNet_b0(nn.Module):
    def __init__(self):
        super(EfficientNet_b0, self).__init__()
        self.model = efficientnet_pytorch.EfficientNet.from_pretrained('efficientnet-b0')
        
        self.classifier_layer = nn.Sequential(
            nn.Linear(1280 , 512),
            nn.BatchNorm1d(512),
            nn.Dropout(0.2),
            nn.Linear(512 , 256),
            nn.Linear(256 , 104)
        )
        
    def forward(self, inputs):
        x = self.model.extract_features(inputs)

        # Pooling and final linear layer
        x = self.model._avg_pooling(x)
        x = x.flatten(start_dim=1)
        x = self.model._dropout(x)
        x = self.classifier_layer(x)
        return x
    
model = EfficientNet_b0()
model = model.to(device)