@ptrblck I created 3 folder for train, val and test and randomly select 60% of image and keep it in train, 20% for val and 20% test.
How to disable dropout layers?
This is my code
feature_extract = True
sm = nn.Softmax(dim = 1)
def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = True    
def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):
    if model_name == "alexnet":
        """ Alexnet
        """
        model= models.alexnet(pretrained=use_pretrained)
        set_parameter_requires_grad(model, feature_extract)
        num_ftrs = model.classifier[6].in_features
        model.classifier[6] = nn.Linear(num_ftrs,num_classes)
        input_size = 224
        

    return model
model = initialize_model(model_name, num_classes, feature_extract, use_pretrained=True)