ML engineers aren’t software engineers who happen to write some Python. They live in notebooks, build training loops, fight CUDA dependency hell, and write code that often exists in a Jupyter cell for six months before it becomes a real file. The AI coding tools built for web and backend devs have caught up with this workflow—some of them, at least.

Here’s who does it well, who’s faking it, and one mistake every AI tool makes with your PyTorch setup.

The CUDA trap — the one mistake every AI coding tool makes

Before comparing features, there’s a specific failure mode you need to know about. Ask any AI coding assistant to help you set up PyTorch with GPU support and there’s a good chance it gives you:

pip install torch torchvision torchaudio

That installs the CPU-only version. Your code will run—silently—on CPU. Training a job that should take 3 minutes takes 3 hours. torch.cuda.is_available() returns False and you spend an hour debugging what you thought was a data pipeline issue.

The correct command for CUDA 12.4:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

PyTorch 2.7—the current stable release—supports CUDA 12.4, 12.6, and 12.8. The 12.6 and 12.8 wheels are now the default pre-compiled binaries. Which one you need depends on your driver. Run nvidia-smi (shows the maximum CUDA version your driver supports) and nvcc --version (shows your installed toolkit version). The driver CUDA version must be ≥ the runtime CUDA version you compile against.

None of the mainstream AI coding tools know your GPU config by default. The fix: paste your nvidia-smi output and nvcc --version into chat before asking any CUDA installation or configuration questions. Every tool on this list gets it right with that context. The failure is the default assumption, not the reasoning capability.

If you’re evaluating GPUs for local ML training, runaihome.com’s GPU buying guide for local AI covers VRAM requirements and CUDA compatibility by GPU generation.

GitHub Copilot — best Jupyter integration for most ML engineers

Pricing: Free (2,000 completions/mo), Pro $10/mo, Pro+ $39/mo, Business $19/user/mo, Enterprise $39/user/mo.

Copilot has the most mature Jupyter integration of any general-purpose AI coding tool. In VS Code, it works inside notebook cells without workarounds—autocomplete, inline edits, and Copilot Chat all understand cell context. Ask it to refactor a data loading cell and it modifies the right cell, not a nearby one.

For JupyterLab users (the native browser interface), Notebook Intelligence (NBI) solves the gap. NBI is a free, open-source JupyterLab extension that connects Copilot, Ollama, Claude Code, or any OpenAI-compatible endpoint directly into JupyterLab. You get chat, inline edits, autocomplete, and a full notebook agent—inside JupyterLab, not VS Code. Install with:

pip install notebook-intelligence

If you’re on Copilot Free or Pro, there’s no additional cost. NBI uses whatever LLM provider you configure.

Jupyter AI v3.0 (released April 1, 2026) is the other free path. It’s an official JupyterLab extension from the Project Jupyter organization, now supporting 1,000+ LLMs via LiteLLM and ACP agents including Claude, Codex, Gemini, Kiro, and OpenCode. The v3 architecture lets you configure custom MCP servers in .jupyter/mcp_settings.json, so ML-specific tools (a vector database, a HuggingFace model index) can be wired into notebook chat. For engineers who want fine-grained control over model routing—fast model for autocomplete, capable model for architecture questions—Jupyter AI v3 with BYOK is worth the setup time.

What Copilot still lacks: cross-cell semantic awareness. It understands individual cells well but struggles to reason about state that spans multiple cells (a variable defined 8 cells earlier that a new cell should reference). For deeply stateful exploration notebooks—which is most ML research work—this is a real gap. The practical workaround is to keep context tight: fresh kernel, clear variable names, and don’t let your notebook grow past 40 cells before refactoring into a .py script.

Copilot also now supports Hugging Face inference providers directly in Copilot Chat inside VS Code. You can switch to specialized models trained on HuggingFace documentation when debugging transformers code—useful when the default model gives you syntax from a deprecated API version.

JetBrains DataSpell 2026.1 — the IDE built for data scientists

Pricing: From $229/year (individual, first year); JetBrains AI Pro bundled at no additional cost since the 2026.1 release (March 2026).

DataSpell exists specifically for data science and ML work. It’s a JetBrains IDE built on the same engine as PyCharm, but with Jupyter notebooks as a first-class citizen rather than an extension. The distinction matters: cells are rendered natively, kernel management is built in, and the AI chat understands cell boundaries without the ambiguity you get when squeezing Jupyter into a generic IDE.

The 2026.1 release added two things ML engineers care about:

AI Agents Ecosystem via ACP Registry. DataSpell’s AI chat can now connect to Claude Agent, Codex, Cursor, and any ACP-compatible agent. The registry lets you discover and install agents with one click from inside the IDE. For teams doing multi-agent work—routing code generation to Sonnet, code review to Opus, and data quality analysis to a specialized model—this is now built into the IDE rather than hacked together externally.

AI Pro bundled. Every paid DataSpell subscription now includes JetBrains AI Pro, which costs $10/month as a standalone add-on. AI Pro gives you advanced completions, AI chat with model selection (GPT-5, Claude, Gemini), and data-specific features: Visualization Cells (added in 2025.3) integrate with AI chat so you can ask “what’s causing the bimodal distribution in this histogram?” and get an analysis. Automated data quality insights flag anomalies in your data exploration cells.

The 2026.1 release also added PDF export for notebooks—a small feature but one that academic ML researchers and anyone doing client presentations will appreciate.

Where DataSpell loses: multi-file Python pipeline work. If your training script spans a dataset.py, model.py, trainer.py, and config.yaml, DataSpell’s cross-file context isn’t as strong as Cursor’s. It’s an IDE optimized for data exploration and analysis. The boundary where it starts to underperform is roughly when your project has more .py files than .ipynb files.

For ML engineers at companies on JetBrains All Products Pack licenses, DataSpell is included. At ~$229/year standalone with AI Pro bundled, the value is strong for notebook-heavy work.

Cursor 1.0 — best for .py ML pipelines, notebooks still catching up

Pricing: Free (2K fast requests/mo), Pro $20/mo, Pro+ $60/mo, Ultra $200/mo, Teams $40/user/mo.

Cursor 1.0 (released early 2026) officially added Jupyter notebook support. The agent can create and edit multiple cells inside .ipynb files—but only with Sonnet models, and with real limitations that matter for ML work.

The core problem is structural: .ipynb files are JSON. When the Cursor agent edits a notebook, it’s rewriting a JSON document, not a Python file. Cell execution state isn’t tracked by the agent, so it can’t know what df holds after 6 cells of transforms, or that model was defined three cells up. The result: suggestions that are syntactically correct but semantically misaligned with your notebook’s current state.

Current best practices for Cursor + ML work:

Use .py files with # %% cell markers for training scripts, data pipelines, and anything you’ll run repeatedly. Cursor’s multi-file context handles these better than any other tool on this list. When you’re refactoring a monolithic train.py into a proper CLI-driven setup, or migrating from PyTorch Lightning to raw PyTorch across multiple files, Cursor’s agent mode is the fastest option available.

For actual .ipynb editing, the community-maintained cursor-notebook-mcp server (GitHub: jbeno/cursor-notebook-mcp) provides a cleaner path. It lets the Cursor agent interact with notebook cells via the Model Context Protocol and execute them directly—working around the JSON fragmentation problem.

Reserve native .ipynb notebooks in Cursor for light exploration. Don’t route heavy multi-cell rewriting through Cursor’s agent in a raw notebook file.

Where Cursor is irreplaceable: complex ML pipelines. A production ML project typically has a training script, dataset class, model definition, config file, and evaluation scripts. Cursor’s cross-file understanding and agent mode handle this structure better than anything else. When your data augmentation pipeline has a bug that’s caused by an interaction between your dataset class and your training loop, Cursor finds it. Copilot and DataSpell are weaker here.

Cursor 1.0 also shipped BugBot—automated PR code review that’s now at 78% resolution rate for issues it flags. For ML engineers who write code reviews on model changes or data preprocessing updates, that’s a meaningful catch rate.

See our Cursor IDE review for full pricing details and feature breakdown.

Google Colab AI — when you’re training in the cloud anyway

Pricing: Free, Pro $9.99/mo, Pro+ $49.99/mo.

If you’re already training on Colab (for the free T4 GPUs or because your institution provides Pro credits), the built-in AI is worth using. Free tier users have access to Gemini models via the google.colab.ai Python library—specifically gemini-2.5-flash and gemini-2.5-flash-lite within monthly limits.

The AI Prompt Cell is the standout: a dedicated cell type where you describe what you want in natural language and Colab generates the Python code. For standard ML patterns—training loop skeleton, checkpoint saving and loading, learning rate scheduling with warmup—the Gemini models handle these well.

Colab AI has one structural advantage none of the other tools have: it knows your runtime environment automatically. Ask it to install PyTorch and it correctly generates the install command for the Colab GPU runtime’s CUDA version, bypassing the CUDA trap entirely. It’s the only tool on this list that solves that problem without you having to inject context.

The ceiling is clear: Colab AI doesn’t follow you outside Colab. When your model needs to go to production—a FastAPI serving endpoint, a model registry push, a Docker container—Colab AI stops being useful. It’s a tool for the exploration and training phase only.

For cloud GPU alternatives to Colab, RunPod lets you spin up a persistent machine with your exact CUDA and framework versions, which is useful when Colab’s ephemeral runtimes break your training runs at the 3-hour mark.

Comparison table

Tool	Jupyter .ipynb editing	PyTorch / CUDA awareness	Multi-file pipeline work	HuggingFace support	Individual monthly cost
Copilot Pro + NBI	Native in VS Code; NBI for JupyterLab	Strong pattern recognition	Good	HF provider in Chat	$10
DataSpell + AI Pro	Best integrated native experience	Good, data-specific insights	Moderate	Good	~$19 (bundled)
Cursor Pro	Limited (Sonnet only; MCP workaround)	Strong with context injection	Best	Strong	$20
Jupyter AI v3	Native JupyterLab	Depends on model choice	N/A	Via LiteLLM	Free + model costs
Colab AI (Pro)	Native (Colab only)	Automatic CUDA match	N/A	Limited	$9.99

Five things AI tools still get wrong with ML code

1. CUDA installation (covered in the opening—fix: inject nvidia-smi and nvcc --version first)

2. Deprecated PyTorch APIs. Tools with training data before mid-2025 suggest torch.cuda.amp.autocast(), which is deprecated in PyTorch 2.x. The correct form is torch.amp.autocast("cuda"). This compiles and runs without a warning if you’re not checking deprecation logs—it just silently works worse in some edge cases.

3. HuggingFace version splits. The transformers library had breaking changes between 4.x and 5.x, particularly in the tokenizer API. AI tools often mix both versions in the same generated code block. Always tell the tool your installed version: pip show transformers | grep Version.

4. DataLoader workers on macOS. AI tools routinely suggest num_workers=4 or higher. On Apple Silicon with the MPS backend, this causes “Too many open files” errors. The correct default for MPS is num_workers=0. DataSpell’s macOS-aware context sometimes catches this; Cursor and Copilot usually don’t.

5. Mixed precision on consumer GPUs. Suggesting torch.float16 on RTX 3070/3080/4070 ignores that bfloat16 is better supported on Ampere and Ada Lovelace consumer hardware. float16 on these GPUs requires gradient scaling to avoid underflow; bfloat16 doesn’t. The fix is to specify dtype=torch.bfloat16 for any Ampere/Ada RTX card.

The universal fix for most of these: dump your environment state into chat before asking ML configuration questions. Running pip list | grep -E 'torch|transformers|accelerate' and python -c "import torch; print(torch.__version__, torch.version.cuda)" takes 10 seconds and prevents 90% of version-mismatch errors.

Honest take

DataSpell 2026.1 is the right choice if notebooks are 80%+ of your day. AI Pro is now bundled, the integration quality for Jupyter is the best in the market, and the ACP registry adds the model flexibility that was the main reason to look elsewhere.

Copilot Pro at $10/month is the pragmatic choice for everyone else—especially anyone already in VS Code. Add NBI for JupyterLab sessions and you have solid Jupyter coverage at the lowest price point on this list.

Cursor Pro at $20/month is the tool you add for serious Python pipeline work, not the one you use as your notebook editor. If your ML work is split between exploration in notebooks and production-quality training pipelines, the $30/month combination of Copilot Pro + Cursor Pro is the most capable setup available.

Jupyter AI v3 is worth adding to your JupyterLab installation regardless of what else you’re running—it’s free, it’s notebook-native, and v3’s MCP support lets you wire in domain-specific context that no commercial tool will give you out of the box.

If you’re optimizing cloud GPU spend alongside your ML tool choices, the GPU Cost Optimization Playbook covers CUDA environment setup, spot instance strategies, and model parallelism cost tradeoffs.

Sources

• GitHub Copilot plans and pricing — github.com/features/copilot/plans
• Cursor pricing — cursor.com/pricing
• DataSpell 2026.1: AI Agents Ecosystem, Export Notebooks to PDF — JetBrains Blog
• Jupyter AI v3 — GitHub (jupyterlab/jupyter-ai)
• Notebook Intelligence (NBI) — GitHub (notebook-intelligence/notebook-intelligence)
• PyTorch 2.7 install guide — pytorch.org/get-started/locally/
• Google Colab Pro pricing — colab.research.google.com/signup
• Cursor 1.0 changelog: BugBot, Background Agent, Jupyter support — cursor.com/changelog/1-0
• Hugging Face inference providers for GitHub Copilot Chat — huggingface.co/docs/inference-providers
• JetBrains AI Plans and Pricing — jetbrains.com/ai-ides/buy/

Last updated May 24, 2026. Pricing and features change frequently; verify current state before purchasing.