```
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split
from grokfast import gradfilter_ema  # Import Grokfast's EMA filter
import matplotlib.pyplot as plt
import time
from torch.amp import autocast, GradScaler

# Define the MLP Model
class SimpleMLP(nn.Module):
    def __init__(self):
        super(SimpleMLP, self).__init__()
        self.flatten = nn.Flatten()
        self.model = nn.Sequential(
            nn.Linear(28*28, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    
    def forward(self, x):
        x = self.flatten(x)
        return self.model(x)

# Prepare the MNIST Dataset
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_val_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

train_size = int(0.8 * len(train_val_dataset))
val_size = len(train_val_dataset) - train_size
train_dataset, val_dataset = random_split(train_val_dataset, [train_size, val_size])

# Optimized DataLoader Parameters
batch_size = 1024  # Increased batch size
num_workers = 8    # Increased number of workers based on CPU cores
train_loader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=num_workers,
    pin_memory=True
)
val_loader = DataLoader(
    val_dataset,
    batch_size=batch_size,
    shuffle=False,
    num_workers=num_workers,
    pin_memory=True
)
test_loader = DataLoader(
    test_dataset,
    batch_size=batch_size,
    shuffle=False,
    num_workers=num_workers,
    pin_memory=True
)

# Training Function with Mixed Precision
def train_model(model, optimizer, criterion, train_loader, val_loader, use_grokfast=False, 
               epochs=10, print_every=1, alpha=0.98, lamb=2.0, use_amp=True):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    
    grads = None  # Initialize grads for Grokfast
    metrics = {
        'train_loss': [],
        'val_loss': [],
        'train_acc': [],
        'val_acc': []
    }
    
    scaler = GradScaler() if use_amp else None
    
    for epoch in range(1, epochs + 1):
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        
        start_time = time.time()
        
        for data, target in train_loader:
            data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
            
            optimizer.zero_grad()
            
            with autocast(device_type="cuda", enabled=use_amp):
                outputs = model(data)
                loss = criterion(outputs, target)
            
            if use_amp:
                scaler.scale(loss).backward()
                if use_grokfast:
                    grads = gradfilter_ema(model, grads=grads, alpha=alpha, lamb=lamb)
                scaler.step(optimizer)
                scaler.update()
            else:
                loss.backward()
                if use_grokfast:
                    grads = gradfilter_ema(model, grads=grads, alpha=alpha, lamb=lamb)
                optimizer.step()
            
            running_loss += loss.item() * data.size(0)
            _, predicted = torch.max(outputs.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
        
        train_loss = running_loss / len(train_loader.dataset)
        train_acc = correct / total
        
        # Validation
        model.eval()
        val_loss = 0.0
        correct = 0
        total = 0
        
        with torch.no_grad():
            for data, target in val_loader:
                data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
                with autocast(device_type="cuda", enabled=use_amp):
                    outputs = model(data)
                    loss = criterion(outputs, target)
                
                val_loss += loss.item() * data.size(0)
                _, predicted = torch.max(outputs.data, 1)
                total += target.size(0)
                correct += (predicted == target).sum().item()
        
        val_loss /= len(val_loader.dataset)
        val_acc = correct / total
        
        metrics['train_loss'].append(train_loss)
        metrics['val_loss'].append(val_loss)
        metrics['train_acc'].append(train_acc)
        metrics['val_acc'].append(val_acc)
        
        if epoch % print_every == 0:
            epoch_time = time.time() - start_time
            print(f'Epoch [{epoch}/{epochs}] - Time: {epoch_time:.2f}s')
            print(f'Train Loss: {train_loss:.4f} | Train Acc: {train_acc*100:.2f}%')
            print(f'Val Loss: {val_loss:.4f} | Val Acc: {val_acc*100:.2f}%\n')
    
    return metrics

# Evaluation Function
def evaluate_model(model, test_loader):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()
    
    correct = 0
    total = 0
    
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
            outputs = model(data)
            _, predicted = torch.max(outputs.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    
    test_acc = correct / total
    return test_acc

# Define loss criterion
criterion = nn.CrossEntropyLoss()

# ======= Baseline Model =======
model_baseline = SimpleMLP()
optimizer_baseline = optim.Adam(model_baseline.parameters(), lr=0.001, weight_decay=0.01)

print("Training Baseline Model (Without Grokfast)...\n")
metrics_baseline = train_model(
    model=model_baseline,
    optimizer=optimizer_baseline,
    criterion=criterion,
    train_loader=train_loader,
    val_loader=val_loader,
    use_grokfast=False,
    epochs=100,          # Set to 10 epochs
    print_every=5,      # Print every epoch
    use_amp=True        # Enable mixed precision
)

# ======= Model with Grokfast =======
model_grokfast = SimpleMLP()
optimizer_grokfast = optim.Adam(model_grokfast.parameters(), lr=0.001, weight_decay=0.01)

print("\nTraining Model with Grokfast...\n")
metrics_grokfast = train_model(
    model=model_grokfast,
    optimizer=optimizer_grokfast,
    criterion=criterion,
    train_loader=train_loader,
    val_loader=val_loader,
    use_grokfast=True,
    epochs=100,          # Set to 10 epochs
    print_every=5,      # Print every epoch
    alpha=0.98,
    lamb=2.0,
    use_amp=True        # Enable mixed precision
)

# Evaluate Models on Test Set
test_acc_baseline = evaluate_model(model_baseline, test_loader)
print(f'\nBaseline Model Test Accuracy: {test_acc_baseline*100:.2f}%')

test_acc_grokfast = evaluate_model(model_grokfast, test_loader)
print(f'Grokfast Model Test Accuracy: {test_acc_grokfast*100:.2f}%')

# Plotting the Results
epochs_plot = range(1, 101)

plt.figure(figsize=(12, 10))

# Train Loss
plt.subplot(2, 2, 1)
plt.plot(epochs_plot, metrics_baseline['train_loss'], label='Baseline')
plt.plot(epochs_plot, metrics_grokfast['train_loss'], label='Grokfast')
plt.xlabel('Epoch')
plt.ylabel('Training Loss')
plt.title('Training Loss Comparison')
plt.legend()

# Validation Loss
plt.subplot(2, 2, 2)
plt.plot(epochs_plot, metrics_baseline['val_loss'], label='Baseline')
plt.plot(epochs_plot, metrics_grokfast['val_loss'], label='Grokfast')
plt.xlabel('Epoch')
plt.ylabel('Validation Loss')
plt.title('Validation Loss Comparison')
plt.legend()

# Train Accuracy
plt.subplot(2, 2, 3)
plt.plot(epochs_plot, [acc * 100 for acc in metrics_baseline['train_acc']], label='Baseline')
plt.plot(epochs_plot, [acc * 100 for acc in metrics_grokfast['train_acc']], label='Grokfast')
plt.xlabel('Epoch')
plt.ylabel('Training Accuracy (%)')
plt.title('Training Accuracy Comparison')
plt.legend()

# Validation Accuracy
plt.subplot(2, 2, 4)
plt.plot(epochs_plot, [acc * 100 for acc in metrics_baseline['val_acc']], label='Baseline')
plt.plot(epochs_plot, [acc * 100 for acc in metrics_grokfast['val_acc']], label='Grokfast')
plt.xlabel('Epoch')
plt.ylabel('Validation Accuracy (%)')
plt.title('Validation Accuracy Comparison')
plt.legend()

plt.tight_layout()
plt.savefig('comparison.png')  # Save the plot instead of displaying
print("\nTraining complete. Comparison plot saved as 'comparison.png'.")
```
modify this code to say how long it took for n epochs, instead of per epoch, where n is the print_every variable