Here is AlexNet and train function implementation . class AlexNet_Modified(nn.Module):
    def __init__(self):
        super(AlexNet_Modified, self).__init__()

        self.features = nn.Sequential(
            nn.Conv2d(3, 96, kernel_size=11, stride=4),
            nn.BatchNorm2d(96),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(96, 256, kernel_size=5, padding=2),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(256, 384, kernel_size=3, padding=1),
            nn.BatchNorm2d(384),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 384, kernel_size=3, padding=1),
            nn.BatchNorm2d(384),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(256 * 6 * 6, 4096),
            nn.BatchNorm1d(4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.3),
            nn.Linear(4096, 4096),
            nn.BatchNorm1d(4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, len(folder_names)),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x and def train_model(model, device, train_loader, test_loader, epochs = None, learning_rate = None, optimization_technique = None,  patience=None, scheduler_patience=None, num_batches = None, **kwargs):

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    if optimization_technique == 'learning_rate_scheduler' and scheduler_patience:
        scheduler = ReduceLROnPlateau(optimizer, 'min', patience=scheduler_patience, verbose=True)

    train_losses = []
    train_accuracies = []
    test_losses = []
    test_accuracies = []

    best_loss = float('inf')
    stopping_counter = 0

    for epoch in range(epochs):
        running_train_loss = 0.0
        running_train_acc = 0
        num_samples_train = 0

        for X_batch, y_batch in train_loader:
            X_batch, y_batch = X_batch.to(device, non_blocking=True), y_batch.to(device, non_blocking=True)

            optimizer.zero_grad()
            y_pred = model(X_batch)
            loss = criterion(y_pred, y_batch)
            loss.backward()
            optimizer.step()

            running_train_loss += loss.item()
            running_train_acc += torch.sum(torch.eq(y_pred.argmax(dim=1), y_batch)).item()
            num_samples_train += y_batch.size(0)

        train_losses.append(running_train_loss / len(train_loader))
        train_accuracies.append(running_train_acc / num_samples_train)

        # Testing segment
        running_test_loss = 0.0
        running_test_acc = 0
        num_samples_test = 0

        with torch.no_grad():
            for X_batch, y_batch in test_loader:
                X_batch, y_batch = X_batch.to(device, non_blocking=True), y_batch.to(device, non_blocking=True)

                y_pred = model(X_batch)
                loss = criterion(y_pred, y_batch)

                running_test_loss += loss.item()
                running_test_acc += torch.sum(torch.eq(y_pred.argmax(dim=1), y_batch)).item()
                num_samples_test += y_batch.size(0)

        test_losses.append(running_test_loss / len(test_loader))
        test_accuracies.append(running_test_acc / num_samples_test)

        # Early stopping
        if optimization_technique == 'early_stopping' and patience:
            if test_losses[-1] < best_loss:
                best_loss = test_losses[-1]
                stopping_counter = 0
            else:
                stopping_counter += 1
            if stopping_counter > patience:
                print(f"Early stopping at epoch {epoch+1}/{epochs}")
                break

        # Learning rate scheduler
        if optimization_technique == 'learning_rate_scheduler' and scheduler_patience and scheduler:
            scheduler.step(test_losses[-1])

        print(f"Epoch: {epoch+1}/{epochs}, Training Loss: {train_losses[-1]:.4f}, Test Loss: {test_losses[-1]:.4f}, Training Accuracy: {train_accuracies[-1]: .4f}, Test Accuracy: {test_accuracies[-1]:.4f}")

    return train_losses, train_accuracies, test_losses, test_accuracies.              Now improve the AlextNEt to perform better for a SVNH dataset and also I'm using data augmentation for SVHN task. Can you also implement steps to import the dataset and also use the code to train the modified model on SVHN dataset.