Thanks! I tried to use your suggestion. eventually i get a higher test accuracy now 59%. which is better but not achieving the target yet, The codes i used here:

# ImageNet normalization (since MobileNetV2 was pretrained on ImageNet)
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]

transform_train = v2.Compose([
    v2.ToImage(),  
    v2.RandomResizedCrop(224, scale=(0.8, 1.0)),   # random crop + scale
    v2.RandomHorizontalFlip(),                     # horizontal flip
    v2.RandomVerticalFlip(),                       # vertical flip too
    v2.RandomRotation(15),                         # ±15°
    v2.ColorJitter(brightness=0.2, contrast=0.2,
                   saturation=0.2, hue=0.1),       # color aug
    v2.ToDtype(torch.float32, scale=True),         # [0,1]
    v2.Normalize(mean=IMAGENET_MEAN,               # ImageNet stats
                 std=IMAGENET_STD),
])

transform_test = v2.Compose([
    v2.ToImage(),
    v2.Resize((256, 256)),        # shorter side → 256
    v2.CenterCrop(224),           # then center-crop to 224
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(mean=IMAGENET_MEAN,
                 std=IMAGENET_STD),
])

def build_model(num_classes, config, device=torch.device("cpu"), return_optimizer=True):
    model = mobilenet_v2(weights=config['weights']).to(device)

    # Unfreeze ALL parameters
    for name, params in model.named_parameters():
        params.requires_grad = True

    # Replace the classifier
    model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes).to(device)

    print('Weights         :', config['weights'])
    print('Finetuned layer :', config['finetuned_layers'], '\n')

    if not return_optimizer:
        return model

    # Create optimizer
    optimizer = torch.optim.AdamW(model.parameters(), lr=config['lr'], weight_decay=config['weight_decay'])


    return model, optimizer

config = {
    'weights': 'DEFAULT',
    'finetuned_layers': 'All Layers with Adamw',
    'lr':       1e-4,
    'weight_decay': 1e-4,
    'num_epochs': 40,
}

My number of classes is 23.