PyTorch vs TensorFlow

The framework war is effectively decided: PyTorch won research; TensorFlow keeps a strong production niche. Most new projects pick PyTorch unless there’s a specific reason (TPU, mobile/edge, existing TF infra).

Side-by-side

Code comparison: same model

PyTorch:

import torch
from torch import nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

model = Net().to("cuda")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

for xb, yb in loader:
    xb, yb = xb.to("cuda"), yb.to("cuda")
    logits = model(xb)
    loss = loss_fn(logits, yb)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

TensorFlow / Keras:

from tensorflow import keras
from tensorflow.keras import layers

model = keras.Sequential([
    layers.Dense(128, activation="relu", input_shape=(784,)),
    layers.Dense(10),
])
model.compile(
    optimizer=keras.optimizers.AdamW(1e-3),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
)
model.fit(x_train, y_train, batch_size=64, epochs=10)

Same model, very different feel:

• Keras hides the loop — easier when standard, awkward when you need custom logic.
• PyTorch shows the loop — more boilerplate but full control everywhere.

Where they differ in practice

Debuggability

PyTorch wins. Set a breakpoint anywhere, print tensors mid-forward. TF eager mode helps, but @tf.function-traced code is harder to debug (you’re inside a graph, not Python).

Production deploy

TensorFlow leads on certain paths:

• TFLite for mobile/microcontrollers — most mature solution.
• TensorFlow.js for browser inference — best-in-class.
• TF Serving — battle-tested gRPC model server.

PyTorch has caught up:

• vLLM, TensorRT-LLM, Triton Inference Server — dominant for LLM serving.
• torch.compile matches @tf.function for graph optimization.
• ONNX export for interop with other runtimes.

Distributed / large-scale training

PyTorch dominates for modern LLM training:

• FSDP (Fully Sharded Data Parallel) — built into PyTorch.
• DeepSpeed, Megatron-LM — PyTorch-native.
• accelerate by Hugging Face — wraps DDP/FSDP/DeepSpeed.

TF has tf.distribute.Strategy — mirrored, multi-worker, TPU. Good for TPU specifically; PyTorch/XLA is improving but TF is the canonical TPU framework.

Hardware accelerators

• NVIDIA GPUs — both excellent. PyTorch slightly ahead on bleeding-edge CUDA features.
• TPUs — TF/JAX preferred; PyTorch/XLA is workable but less mature.
• AMD GPUs (ROCm) — PyTorch better support.
• Apple Silicon (MPS) — PyTorch has mps backend; TF via tensorflow-metal.

Pretrained model availability

PyTorch wins by a wide margin. Hugging Face’s transformers, most LLM checkpoints (Llama, Mistral, Qwen, DeepSeek), diffusion models — all PyTorch-first. TF variants exist but lag.

Learning curve

PyTorch is generally considered easier to pick up for Python users — NumPy-like, no graph-mode mental model. TF/Keras is easier for the common case (.fit()) but harder once you go off the rails.

When to pick TensorFlow

• You need TPU training at scale (Google ecosystem).
• You’re deploying to mobile / microcontrollers (TFLite).
• You need browser inference (TensorFlow.js).
• You’re maintaining an existing TFX pipeline.
• Your team already has TF expertise and the model isn’t an LLM.

When to pick PyTorch

• Research / experimentation — flexibility matters more than fit-and-forget.
• LLMs / generative AI — the entire ecosystem (HF, vLLM, DeepSpeed) is PyTorch.
• Custom architectures with non-standard training loops.
• AMD GPUs or Apple Silicon.
• You want HF transformers off the shelf.

When to pick neither

• JAX for high-performance research, TPU-native, functional style (Google DeepMind models, increasingly common).
• scikit-learn for classical ML (random forests, gradient boosting, SVMs) — neither framework is needed.
• XGBoost / LightGBM for tabular data — usually beats deep learning anyway.
• Hugging Face transformers as the layer above PyTorch for pretrained-model use cases.

Keras 3 — does this change anything?

Keras 3 (2023+) is multi-backend: write once with keras.layers.*, run on TF, JAX, or PyTorch. Lowers the cost of framework choice for new code that fits the Keras API. Still: most cutting-edge research and the HF ecosystem live in raw PyTorch, not Keras.

Migration thoughts

• TF → PyTorch — common path for teams modernizing. The model code translates straightforwardly; the harder part is rebuilding training infra and deploy pipelines.
• PyTorch → TF — rare. Usually only when forced into TPU-only environments or legacy TF infra.

Decision shortcut

Doing LLMs / generative AI?
  → PyTorch (HF ecosystem)

Need TPU at scale?
  → TF or JAX

Deploying to mobile / microcontroller?
  → TF (TFLite)

Browser inference?
  → TF.js

Classical ML (tabular)?
  → scikit-learn / XGBoost — skip both DL frameworks

Otherwise?
  → PyTorch (research-friendly, dominant)

Interview angle

• “PyTorch vs TF — which would you pick for a new project?” — PyTorch for most cases, especially LLMs/research. TF if you specifically need TPU, TFLite, TF.js, or have existing TF infrastructure. The two frameworks have largely converged in capability; ecosystem matters more than features now.
• “Why did PyTorch win research?” — define-by-run dynamic graphs (Python control flow just works), debuggability (regular stack traces, breakpoints anywhere), NumPy-like API. TF 1.x’s static graph model was painful enough that researchers fled.
• “Is TF dead?” — no. It’s the second-most-popular framework, dominant in production for non-LLM use cases, owns mobile/edge and TPU, and Keras 3 makes it backend-agnostic. But mindshare for new projects is mostly PyTorch.
• “What’s torch.compile vs @tf.function?” — both JIT-compile model code into optimized graphs for speed. torch.compile is newer (PyTorch 2.0, 2023) and uses Inductor/Triton; @tf.function has been around since TF 2.0. Similar concept, different implementations.