ML Bug Squash: Identify and Fix Common Bugs in a Training Pipeline

Question

## Problem

You are given a machine learning training script with several embedded bugs. Your task is to identify and fix them.

**Bug hunt — find at least 4 issues in this pseudocode:**

```python
def train(model, X_train, y_train, X_test, y_test):
    # Bug 1: Normalization uses test stats
    mean = X_test.mean(axis=0)
    std  = X_test.std(axis=0)
    X_train = (X_train - mean) / std
    X_test  = (X_test  - mean) / std

# Bug 2: Shuffle happens after train/test split (data already split above)
    shuffle_in_place(X_train, y_train)  # should be before split

for epoch in range(100):
        loss = model.forward(X_train, y_train)
        # Bug 3: Gradient not zeroed before backward
        model.backward(loss)
        model.step(lr=0.01)

# Bug 4: Evaluating on training data, not test data
        acc = model.evaluate(X_train, y_train)
        print(f"Epoch {epoch}: acc={acc}")

return model
```

## Follow-ups

1. Why does normalizing using test statistics cause data leakage?
2. What is the effect of not zeroing gradients — in which framework (PyTorch/TF) does this matter most?
3. How would you structure a training loop to prevent these classes of bugs systematically?
4. What automated checks (e.g., assertions, dataset auditing) would you add before training starts?

ML Bug Squash: Identify and Fix Common Bugs in a Training Pipeline

Question Details

Full Details

Topics

More from Stripe