Abhigyan Chakraborty

Installing PyTorch and TensorFlow

Phase 3: High-Performance Computing — Part 2

(Part 9 of series Blueprint for a Modern Research Computing Environment)

Follow me :

LinkedIn    Website    Website

Quick Summary

This article installs PyTorch and TensorFlow, each in its own dedicated Conda environment, using pip so both frameworks bundle their own CUDA and cuDNN libraries. You’ll verify GPU detection for both, benchmark CPU vs GPU speed with a matrix multiplication test, and — if needed — resolve the WSL library-path issue flagged back in Part 9.

Objective

In Part 9, we set up GPU access — and flagged one thing to expect here in Part 10: TensorFlow’s pip-bundled CUDA libraries sometimes need one extra step to be found on WSL. In this part, we’ll install PyTorch and TensorFlow, verify that both can see and use your GPU, and walk through that extra step if you need it.

By the end, you’ll have:

Content

💡 Getting Unstuck (Expand for AI Troubleshooting Prompts)

AI tools (like Claude, Grok, Gemini, or ChatGPT) give the best troubleshooting advice when they have the exact text of the tutorial. To ensure the AI understands what you are trying to build, we will give it the actual file.

1. Download the tutorial file

2. Upload it to your AI

3. Ask for Help Copy and paste this exact prompt into the chat along with your file:

I have attached the markdown file for the tutorial I am following. Please read it so you understand the specific environment I am trying to build.

I need help with the following:

Step [X]: [paste exact step text from the blog]

Command I ran: [paste exact command]

What happened: [paste full output/error — if the terminal was truly blank, say so explicitly]

Please help me troubleshoot and fix this error. You can use your general knowledge to solve the problem, but your solution MUST align with the architecture and tools taught in the attached file. Do not suggest alternative setups that contradict the tutorial. Once fixed, tell me what to do next in the article.

To go deeper on a step before you run a command:

I have attached the Part9.md file. Look at Step [X] and explain exactly what the command does and why we are doing it before I run it.

Think of this series as the roadmap and your AI assistant as your learning companion.

Prerequisites

PyTorch vs TensorFlow — A Brief Overview

Both are open-source deep learning frameworks running on GPU. PyTorch (Meta) is flexible and dominant in research, while TensorFlow (Google) is more structured and common in production.

Why Two Separate Environments?

PyTorch and TensorFlow can conflict regarding CUDA library or dependency versions. Installing each in its own Conda environment keeps them isolated and stable.

Why We Use pip for Both Frameworks

We use pip because it installs the latest versions directly from PyPI and bundles the required CUDA/cuDNN libraries automatically, removing the need for a system-wide CUDA Toolkit.

Section 1 — Installing PyTorch

💡Expand for details

Step 1 — Create a Dedicated Environment

  conda create -n env_pytorch python=3.11 pip
  conda activate env_pytorch

Step 2 — Get the Correct Install Command

Visit pytorch.org/get-started/locally/ and generate your command using the Pip package option and your nvidia-smi reported CUDA version.

Run the generated command (example):

  pip3 install torch torchvision --index-url [https://download.pytorch.org/whl/cu132](https://download.pytorch.org/whl/cu132)

Step 3 — Register the Kernel

  conda install ipykernel
  python -m ipykernel install --user --name env_pytorch --display-name "Python (env_pytorch)"

Step 4 — Verify PyTorch Installation

Verify version: python -c "import torch; print(torch.__version__)". Verify GPU: python -c "import torch; print(torch.cuda.is_available())".

Section 2 — Installing TensorFlow

💡Expand for details

Step 1 — Create a Dedicated Environment

  conda create -n env_tensorflow python=3.11 pip
  conda activate env_tensorflow

Step 2 — Install TensorFlow

Ensure pip is current, then install with GPU support:

  pip install --upgrade pip
  pip install tensorflow[and-cuda]

Step 3 — Confirm the CUDA Packages

Ensure the NVIDIA packages landed:

  pip list | grep -i nvidia

If empty, reinstall using pip install --force-reinstall "tensorflow[and-cuda]" or pin to a specific version if needed.

Step 4 — Register the Kernel

  conda install ipykernel
  python -m ipykernel install --user --name env_tensorflow --display-name "Python (env_tensorflow)"

Step 5 — Verify Installation and WSL Pathing

Verify: python -c "import tensorflow as tf; print(tf.__version__)". GPU Check: python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))".

If empty on WSL or Linux, fix the library path:

  1. Find the location: pip show nvidia-cudnn-cu12.
  2. Identify the nvidia path: /home/<user>/miniconda3/envs/env_tensorflow/lib/python3.11/site-packages/nvidia.
  3. Apply the library path fix and make it permanent for this environment only.** Rather than adding this to ~/.bashrc (which would apply it globally, even to environments that don’t need it), scope it to env_tensorflow using conda’s activation hooks:
  4. Run each of the following line seperately
  mkdir -p $(conda info --base)/envs/env_tensorflow/etc/conda/activate.d

  cat > $(conda info --base)/envs/env_tensorflow/etc/conda/activate.d/env_vars.sh << 'EOF'

  NVIDIA_BASE=/home/<your-user>/miniconda3/envs/env_tensorflow/lib/python3.11/site-packages/nvidia

  export LD_LIBRARY_PATH=$(find "$NVIDIA_BASE" -maxdepth 2 -type d -name lib | tr '\n' ':')$LD_LIBRARY_PATH

  EOF

Then reactivate the environment to pick it up:

  conda deactivate
  conda activate env_tensorflow

From now on, every time you activate env_tensorflow, this path is set automatically — no need to repeat the fix.

If you’re still seeing an empty list after this:

  1. Confirm nvidia-smi still runs cleanly and shows your GPU.
  2. Check ~/.bashrc for a leftover LD_LIBRARY_PATH=/usr/local/cuda/lib64 line from an older system CUDA Toolkit install — this can conflict with the path you just set.
  3. Use the ChatAI prompt at the top of this article — paste your nvidia-smi output and the output of the GPU detection command — to troubleshoot further.

Section 3 — Verifying GPU Acceleration

Create and run these files to compare CPU vs GPU performance.

PyTorch Benchmark (pytorch_project.py):

💡Expand for details

In Terminal:

  conda activate env_pytorch
  nano pytorch_project.py

Paste the following:

  import torch
  import time
  size = 10000
  # CPU
  a_cpu = torch.randn(size, size)
  b_cpu = torch.randn(size, size)
  start = time.time()
  c_cpu = torch.matmul(a_cpu, b_cpu)
  cpu_time = time.time() - start
  print(f"CPU time: {cpu_time:.4f} seconds")
  # GPU
  a_gpu = a_cpu.cuda()
  b_gpu = b_cpu.cuda()
  torch.cuda.synchronize()
  start = time.time()
  c_gpu = torch.matmul(a_gpu, b_gpu)
  torch.cuda.synchronize()[]
  gpu_time = time.time() - start
  print(f"GPU time: {gpu_time:.4f} seconds")
  print(f"Speedup: {cpu_time / gpu_time:.1f}x")

Then (ctrl+X, Y, Enter). Run the code:

  python pytorch_project.py

TensorFlow Benchmark (tensorflow_project.py):

💡Expand for details

In Terminal:

  conda activate env_tensorflow
  nano tensorflow_project.py

Paste the following:

  import tensorflow as tf
  import time
  size = 10000
  # CPU
  with tf.device('/CPU:0'):
      a_cpu = tf.random.normal([size, size])
      b_cpu = tf.random.normal([size, size])
      start = time.time()
      c_cpu = tf.matmul(a_cpu, b_cpu)
      cpu_time = time.time() - start
      print(f"CPU time: {cpu_time:.4f} seconds")
  # GPU
  with tf.device('/GPU:0'):
      a_gpu = tf.random.normal([size, size])
      b_gpu = tf.random.normal([size, size])
      start = time.time()
      c_gpu = tf.matmul(a_gpu, b_gpu)
      gpu_time = time.time() - start
      print(f"GPU time: {gpu_time:.4f} seconds")
      print(f"Speedup: {cpu_time / gpu_time:.1f}x")

Then (ctrl+X, Y, Enter). Run the code:

  python tensorflow_project.py

CPU-Only Users

If you lack an NVIDIA GPU, install the CPU-only version:

What’s Next

What You’ve Done:

Next: Part 10 — Building Reproducible Research Workflows | Previous: Part 8 — Enabling GPU Computing in WSL2 and Linux with CUDA

All Blogs

This site is open source. Improve this page.