How to Debug Python in VS Code for Machine Learning Projects

Machine learning code fails in ways that are uniquely frustrating. A model trains for six hours and silently produces garbage predictions. A tensor shape mismatch throws an error on line 247 of a training loop. A data pipeline leaks memory so slowly you don’t notice until your cloud bill arrives. These aren’t the kinds of bugs you fix with a quick print() statement — and yet many ML practitioners rely on exactly that. If you’ve been logging values to the console and squinting at terminal output, this guide is for you.

VS Code has one of the most capable Python debugging environments available, and when you know how to configure it specifically for machine learning workflows — notebooks, training loops, GPU pipelines — it becomes a genuinely transformative tool. This guide walks you through everything that matters: setup, configuration, techniques, and the specific debugging patterns that apply to ML projects in particular.

Getting Your Environment Ready

Before you touch the debugger, your environment needs to be correct. The most common reason the VS Code debugger behaves unexpectedly is a mismatch between the Python interpreter it’s using and the one your project actually runs on.

Open the Command Palette (Ctrl+Shift+P on Windows/Linux, Cmd+Shift+P on Mac) and type Python: Select Interpreter. Choose the interpreter from your virtual environment — whether that’s a venv, conda environment, or a Docker-bound interpreter. If you’re using conda, the interpreter will typically be at something like ~/anaconda3/envs/my_ml_env/bin/python.

You’ll also want these VS Code extensions installed:

  • Python (by Microsoft) — the core extension, non-negotiable
  • Pylance — language server for intelligent type checking and autocomplete
  • Jupyter — essential for debugging notebook cells interactively
  • Python Debugger (Debugpy) — ships with the Python extension but worth knowing by name

Once these are installed and your interpreter is selected, you’re ready to configure the debugger properly.

Understanding launch.json for ML Projects

The launch.json file in the .vscode/ folder is where you define how the debugger starts your program. A generic Python configuration works fine for simple scripts, but ML projects usually need something more specific.

Here’s a well-structured launch.json for a machine learning project:

{
  "version": "0.2.0",
  "configurations": [
    {
      "name": "Train Model",
      "type": "debugpy",
      "request": "launch",
      "program": "${workspaceFolder}/train.py",
      "args": ["--epochs", "10", "--batch-size", "32", "--lr", "0.001"],
      "console": "integratedTerminal",
      "env": {
        "CUDA_VISIBLE_DEVICES": "0",
        "PYTHONPATH": "${workspaceFolder}"
      },
      "justMyCode": false
    },
    {
      "name": "Debug Data Pipeline",
      "type": "debugpy",
      "request": "launch",
      "program": "${workspaceFolder}/data/preprocess.py",
      "console": "integratedTerminal",
      "justMyCode": false
    }
  ]
}

A few things to note here. The args array lets you pass CLI arguments to your training script directly from the debug config — so you can set a small --epochs 2 value to run a fast debug loop. The "justMyCode": false flag is critical: with it set to true (the default), VS Code won’t step into library code like PyTorch or NumPy, which is usually fine. But when a bug lives inside a library call, you need this off to trace it.

The env block lets you set environment variables per-configuration, which is useful for controlling GPU visibility or setting PYTHONPATH so your imports resolve correctly.

Breakpoints: Beyond the Basics

Everyone knows you can click in the gutter to set a breakpoint. What most people don’t use are the conditional breakpoints and logpoints — and for ML debugging, these are far more valuable than standard breakpoints.

Conditional Breakpoints pause execution only when a condition is true. Right-click on a breakpoint and select “Edit Breakpoint” to add a condition. For a training loop, this means you can write:

loss > 10.0

…and the debugger will pause only when your loss explodes — not on every iteration of a 50,000-step training run. You can also use hit counts: pause only on every 500th iteration to sample behavior over time.

Logpoints are even more elegant. Instead of pausing execution, they log a message to the Debug Console when hit. Right-click in the gutter and choose “Add Logpoint”. You can write expressions like:

Epoch {epoch}, loss={loss.item():.4f}, lr={optimizer.param_groups[0]['lr']}

This gives you a clean stream of formatted output without touching your source code at all — and without the performance overhead of re-running the whole script when you want to add or remove a print statement.

Debugging Jupyter Notebooks in VS Code

Notebooks get special treatment. With the Jupyter extension installed, you can set breakpoints directly in notebook cells just like you would in a .py file. Click the gutter to the left of a cell line number, then run the cell in debug mode using the dropdown next to the Run button.

One pattern that works especially well for ML notebooks is to convert exploratory cells into a proper debug session once you’ve isolated a problem. If a cell is throwing a shape mismatch error, right-click on it and use “Debug Cell” — this runs just that cell with the current kernel state intact, pausing at your breakpoint so you can inspect tensors.

The Variables panel in VS Code becomes your best friend during notebook debugging. It shows all variables in the current kernel scope, and for numpy arrays and PyTorch tensors, it renders their dtype, shape, and value summary inline. You don’t have to print anything — just look at the panel.

One important caveat: when debugging notebooks that use GPU tensors, calling .item() or .cpu() in the debug console on a CUDA tensor works fine, but directly printing a large tensor from the Watch panel can cause a slowdown. Be selective about what you evaluate in the watch expressions.

The Debug Console and Watch Panel for Tensor Inspection

When execution is paused at a breakpoint, the Debug Console at the bottom of VS Code becomes an interactive Python REPL with full access to the current scope. This is enormously powerful for ML debugging.

Some practical commands to run in the Debug Console during a paused ML session:

# Check tensor shape and dtype
x.shape, x.dtype

# Check for NaNs (common silent killer in training)
import torch; torch.isnan(x).any()

# Inspect model parameter gradients
[(name, p.grad.norm().item()) for name, p in model.named_parameters() if p.grad is not None]

# Check GPU memory
import torch; torch.cuda.memory_summary()

# View a batch of data as a summary
X_batch.min(), X_batch.max(), X_batch.mean()

The Watch panel lets you pin expressions that re-evaluate every time execution pauses. For a training loop, you might pin loss.item(), model.training, and optimizer.param_groups[0]['lr'] so you can monitor them at every breakpoint hit without typing them again.

Visual: Common ML Bug Types and Where to Catch Them

Where ML Bugs Live — and How the Debugger Catches Them

🔴 Shape Mismatch
Tensor dimensions don’t align in matrix ops or layer inputs.
Breakpoint → Debug Console
x.shape, y.shape
🟠 NaN / Inf Loss
Exploding gradients or bad preprocessing produce silent NaN values.
Conditional BP: loss.isnan()
Watch: loss.item()
🟡 Data Leakage
Train-set statistics bleed into validation transforms or labels.
Step into Dataset.__getitem__
Inspect scaler.mean_
🟢 Device Mismatch
Model is on GPU, input tensor is on CPU — or vice versa.
Watch: x.device
Watch: next(model.parameters()).device

Debugging Training Loops: A Step-by-Step Approach

Training loops are where most ML bugs hide. The loop runs for hours, a metric slowly drifts wrong, and there’s no obvious stack trace. Here’s a disciplined approach to debugging one:

Step 1 — Run one iteration first. Add break after the first batch in your loop and use the debugger to step through the entire forward pass, loss computation, backward pass, and optimizer step. Verify that loss decreases, gradients exist, and weights change.

Step 2 — Use the Call Stack panel. When paused inside a nested function call during your forward pass, the Call Stack panel shows you the exact path of execution. For custom PyTorch modules, this is invaluable — you can click any frame in the stack to jump to that scope and inspect its local variables.

Step 3 — Step into library code when needed. If your loss function is behaving unexpectedly, set "justMyCode": false and step into torch.nn.CrossEntropyLoss or whatever you’re using. You’ll see exactly what intermediate values are computed.

Step 4 — Use exception breakpoints. In the Breakpoints panel, enable “Raised Exceptions” or “Uncaught Exceptions.” This tells the debugger to pause the moment any exception occurs — even if it’s caught internally by a try/except — so you can inspect the live state that caused it.

Step 5 — Validate your data loader separately. Isolate the DataLoader and iterate a few batches in a debug session. Check label distributions, tensor ranges, and whether augmentation is applied consistently.

Remote Debugging for GPU Servers and Docker Containers

Many ML practitioners train on remote GPU servers or inside Docker containers. VS Code handles this well with its Remote – SSH and Dev Containers extensions.

With Remote – SSH, you connect to your training server directly in VS Code — the editor runs locally, but all files, terminals, and debugger processes run on the remote machine. Your launch.json configurations work identically. The debugger attaches over SSH without any extra configuration.

For Docker, add the Dev Containers extension and open your project inside the container. VS Code installs its server components inside the container automatically. You can debug your PyTorch training script with GPU access from within the container, with the full VS Code debugger attached.

If you need to attach the debugger to an already-running Python process — for example, a long-running training job you started without the debugger — use debugpy directly:

import debugpy
debugpy.listen(("0.0.0.0", 5678))
debugpy.wait_for_client()  # pauses execution until debugger connects

Then in your launch.json, add a configuration with "request": "attach" and the appropriate host and port.

Visual: Debugging Workflow Flowchart

ML Debugging Decision Flow in VS Code

Bug detected (wrong output / crash / slow training)
Is it a crash with a traceback?
YES → Enable Exception
Breakpoints, re-run
NO → Use Conditional BP
on metric or loss value
Paused at breakpoint → inspect Variables, Call Stack, Watch panel
Use Debug Console to test fix hypotheses interactively
Apply fix → remove breakpoints → full training run

Common Pitfalls and How to Avoid Them

Even with the debugger configured correctly, a few common mistakes trip people up in ML projects specifically:

Forgetting model.eval() during debugging. If you step through inference code without calling model.eval(), dropout layers remain active and batch normalization uses running statistics differently. This can make your debug session produce different behavior than production. Always check the model.training flag in your Watch panel.

Stepping through multiprocessing DataLoaders. If your DataLoader uses num_workers > 0, breakpoints inside __getitem__ won’t trigger — worker processes are separate from the main debugger process. Set num_workers=0 during debug sessions.

Large tensor evaluation freezing the debugger. Hovering over a large tensor in the Variables panel causes VS Code to evaluate and display it, which can take seconds or freeze the UI entirely. Avoid evaluating tensors larger than a few hundred elements in the watch/hover display; use the Debug Console instead where you control what’s computed.

Not using a small data subset for debugging. Always have a --debug flag in your training script that loads only 100 samples and runs 2 epochs. Attaching the debugger to a full training run on a million samples is painful. Make fast iteration easy.

Conclusion

Debugging machine learning code in VS Code stops being painful once you go beyond the defaults. Conditional breakpoints eliminate the noise of stepping through thousands of loop iterations. The Debug Console gives you an interactive lens into your model’s live state. Proper launch.json configuration means you’re always debugging the right interpreter, with the right environment variables, passing the right arguments.

The shift from print-based debugging to a full debugger session isn’t just about convenience — it’s about the quality of insight you get. When you can pause mid-training loop, inspect every tensor, walk the call stack, and test a fix hypothesis interactively without restarting, you solve bugs faster and understand your models more deeply. That’s a compounding advantage across every project you work on.

Leave a Comment