This article was originally published on aicoderscope.com
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 ha
Top comments (0)