适配CUDA 12.6的PyTorch安装及Jupyter Notebook内核重启问题求助

Hey there, let's work through your PyTorch + CUDA 12.6 setup and Jupyter kernel restart issues together—this is a common scenario, so we’ve got plenty of practical fixes to try out!

一、安装适配CUDA 12.6的PyTorch #

Since PyTorch’s official stable builds currently support up to CUDA 12.4, we’ve got two reliable paths to get it working smoothly with CUDA 12.6:

方法1：利用CUDA向后兼容性安装官方CUDA 12.4版本 #

CUDA is built with backward compatibility in mind, so the CUDA 12.4 PyTorch build will run perfectly with your system’s CUDA 12.6. Here’s how to set it up:

1. First, double-check that your CUDA 12.6 is properly installed and environment variables are configured (make sure PATH and LD_LIBRARY_PATH point to your CUDA 12.6 installation directory).
2. Run the official installation command for CUDA 12.4:

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

3. Verify the setup by opening a Python terminal and running:

   import torch
   print(torch.cuda.is_available())  # Should return True if everything works
   print(torch.version.cuda)         # Will show 12.4, but it uses your system's CUDA 12.6 under the hood
   

方法2：安装PyTorch Nightly版本（原生支持CUDA 12.6） #

If you want a build that explicitly targets CUDA 12.6, go with the nightly pre-release builds—these are updated frequently and already support the latest CUDA versions:

pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu126

Verify with the same Python code above; this time torch.version.cuda will directly show 12.6.

二、解决Jupyter Notebook内核重启问题 #

Kernel restarts usually come from environment mismatches, memory shortages, or dependency conflicts. Let’s try these fixes one by one:

• Fix 1: Install Jupyter directly in your PyTorch environment
  A common mistake is using a global Jupyter installation with a virtual environment’s PyTorch. Instead, activate your PyTorch environment first, then install Jupyter there:

  # Activate your virtual environment (e.g., conda activate your_pytorch_env)
  pip install jupyter notebook
  

  Launch Jupyter from this environment—this ensures the kernel uses the exact same PyTorch setup you installed, eliminating cross-environment conflicts.

• Fix 2: Check for memory overflow
  Kernel restarts often happen when your GPU/CPU runs out of memory. Try these steps:

  • Close any other resource-heavy programs running in the background.
  • Reduce batch sizes or simplify your model architecture in code.
  • Add these lines to the start of your notebook to monitor GPU memory usage:

    import torch
    print("Allocated GPU memory:", torch.cuda.memory_allocated())
    print("Reserved GPU memory:", torch.cuda.memory_reserved())
    

• Fix 3: Reinstall and re-link your environment’s kernel
  If Jupyter isn’t properly recognizing your PyTorch environment, re-setup the kernel:

  # Activate your PyTorch environment
  pip install ipykernel
  python -m ipykernel install --user --name=your_env_name --display-name="PyTorch (CUDA 12.6)"
  

  When you launch Jupyter, select this newly added kernel from the "Kernel > Change Kernel" menu.

• Fix 4: Update all dependent packages
  Outdated packages can cause hidden compatibility issues. Run these commands to update everything to the latest versions:

  pip install --upgrade pip setuptools wheel
  pip install --upgrade torch torchvision torchaudio jupyter ipykernel
  

备注：内容来源于stack exchange，提问作者Srivathsan SK