Torch.compile is decreasing training speed instead increasing it

### 🐛 Describe the bug

Hi, everyone
I turned into Linux to can work with torch.compile as windows doesn't support to test how much will torch.compile helping me speed my training, but the result was disappointed.

that's my code

# Enable TensorFloat32 for better performance
torch.set_float32_matmul_precision('high')

# Enable cuDNN benchmarking for optimal performance
torch.backends.cudnn.benchmark = True

batch_size = 32

train_data = torchvision.datasets.ImageFolder(root=train_dataset_path,
transform=train_transforms,
target_transform=None)

test_data = torchvision.datasets.ImageFolder(root=test_dataset_path,
transform=test_transforms,
target_transform=None)

class_names = train_data.classes

print(len(train_data)), print(len(test_data))

# Create DataLoader for preloaded data
train_loader = DataLoader(
dataset=train_data,
batch_size=batch_size,
num_workers=5,
pin_memory=True,
persistent_workers=True,
prefetch_factor=3,
shuffle=True
)

test_loader = DataLoader(
dataset=test_data,
batch_size=batch_size,
num_workers=5,
pin_memory=True,
persistent_workers=True,
prefetch_factor=3,
shuffle=False
)

def set_device():
if torch.cuda.is_available():
dev = "cuda"
else:
dev = "cpu"
return torch.device(dev)

device = set_device()

scaler = torch.amp.GradScaler('cuda')

def train_nn(model, train_loader, test_loader, criterion, optimizer, n_epochs, is_compiled=False) -> Dict[str, List[float]]:

device = set_device()
best_acc = 0

# Initialize the learning rate scheduler
scheduler = ReduceLROnPlateau(
optimizer, mode='min', factor=0.5, patience=3)

early_stopping = EarlyStopping(patience=11, min_delta=0.001)

results = {
'train_loss': [],
'train_acc': [],
'test_loss': [],
'test_acc': []
}

for epoch in range(n_epochs):

print("Epoch number %d" % (epoch + 1))
start_time = timer()
model.train()

running_loss = 0.0
running_correct = 0.0
total = 0

for data in (train_loader):

images, labels = data
images = images.to(device)
labels = labels.to(device)
total += labels.size(0)

optimizer.zero_grad(set_to_none=True)

# Use AMP for forward and backward pass
with torch.amp.autocast(device_type='cuda'):
outputs = model(images)
loss = criterion(outputs, labels)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

running_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
running_correct += (labels == predicted).sum().item()

epoch_loss = running_loss / len(train_loader)
epoch_acc = 100.00 * running_correct / total

print(" - Training Dataset. Got %d out of %d images correctly (%.3f%%). Epoch loss: %.3f"
% (running_correct, total, epoch_acc, epoch_loss))

test_dataset_acc, test_loss, all_preds, all_labels = evaluate_model_on_test_set(
model, test_loader, criterion)
end_time = timer()

elapsed_time = end_time - start_time
print('Execution time:', time.strftime(
"%H:%M:%S", time.gmtime(elapsed_time)))

print(
f"Checking epoch_loss {epoch_loss} test_loss {test_loss} epoch_acc {epoch_acc} test_dataset_acc {test_dataset_acc}")

if (epoch_loss <= 0.0001 and test_loss <= 0.0001 and
epoch_acc >= 99.999 and test_dataset_acc >= 99.999):
best_acc = test_dataset_acc
save_checkpoint(model, epoch, optimizer, best_acc, is_compiled)
print("Perfect performance achieved! Training stopped.")
break

if test_dataset_acc > best_acc or (test_dataset_acc == best_acc and epoch_loss < results['train_loss'][-1] if results['train_loss'] else True):
best_acc = test_dataset_acc
save_checkpoint(model, epoch, optimizer, best_acc, is_compiled)

results['train_loss'].append(epoch_loss)
results['train_acc'].append(epoch_acc)
results['test_loss'].append(test_loss)
results['test_acc'].append(test_dataset_acc)

scheduler.step(test_loss)
print(f"Current Learning Rate: {optimizer.param_groups[0]['lr']}")

if early_stopping.check_early_stop(test_loss):
print("Early stopping triggered. Training stopped.")
break

print("Training Complete")
return model, results, all_preds, all_labels

EfficientNetV2_S_model = models.efficientnet_v2_s(
weights=models.EfficientNet_V2_S_Weights.IMAGENET1K_V1)

# Modify the first convolutional layer to accept 1-channel (grayscale) input
original_first_conv = EfficientNetV2_S_model.features[0][0]
new_first_conv = nn.Conv2d(
in_channels=1,
out_channels=original_first_conv.out_channels,
kernel_size=original_first_conv.kernel_size,
stride=original_first_conv.stride,
padding=original_first_conv.padding,
bias=False
)

# Initialize the new first layer's weights by averaging the pretrained weights across the RGB channels
with torch.no_grad():
new_first_conv.weight[:, :] = original_first_conv.weight.mean(
dim=1, keepdim=True)

EfficientNetV2_S_model.features[0][0] = new_first_conv

# Modify the classifier's final layer for 4 output classes
num_ftrs = EfficientNetV2_S_model.classifier[1].in_features
num_classes = 4
EfficientNetV2_S_model.classifier[1] = nn.Linear(num_ftrs, num_classes)

EfficientNetV2_S_model = EfficientNetV2_S_model.to(device)

# COMPILE THE MODEL
print("Compiling model with torch.compile...")
EfficientNetV2_S_model = torch.compile(
EfficientNetV2_S_model, mode="max-autotune")
print("Model compilation complete!")

loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(EfficientNetV2_S_model.parameters(),
lr=0.0003, weight_decay=1e-5)

model, results, all_preds, all_labels = train_nn(
EfficientNetV2_S_model, train_loader, test_loader, loss_fn, optimizer, 40, is_compiled)

I hope I'm wrong and help me speeding my training with compile

### Error logs

_No response_

### Versions

Python 3.13.11
torch2.9.1+cuda13.0

cc @jerryzh168 @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @kadeng @amjames @Lucaskabela @jataylo