PyTorch Config Change Reduces Docker Image Size by 78%

Developers at online marketplace giant OLX wrote in a blog post about how they achieved a 78% reduction in Docker image sizes for machine learning services by modifying a single line in their Python requirements configuration.

The optimisation targets the PyTorch library, a dominant AI framework developed by Meta’s Fundamental AI Research lab.

The technical team discovered that default installations of the ‘sentence-transformers’ library automatically pull in extensive NVIDIA CUDA and Triton GPU libraries, even when the service is intended to run on a CPU. 

By adding the flag ‘--extra-index-url https://download.pytorch.org/whl/cpu’ to the requirements.txt file, the image size dropped to 1.75GB.

Himanshu Jhamb, Lead Data Scientist at OLX and the blog’s author, noted that oversized images caused tangible performance issues in production environments. 

He observed that the primary problem “wasn’t the code—it was the Docker image size” and that these 6.5GB of unneeded GPU libraries led to “increased latency during cold starts” and higher system loads during pod restarts. 

This reduction addresses several infrastructure bottlenecks. 

Smaller images reduce the time required for container registries to push and pull data, accelerating CI/CD pipelines and improving autoscaling responsiveness. 

Large images often delay pod readiness, potentially causing traffic spikes when multiple services restart simultaneously.

Further internal analysis showed that the default PyTorch installation allocates approximately 2.8GB to NVIDIA-related files and 913MB to the core torch library. 

In contrast, the CPU-specific build requires only 616MB for the torch component. Jhamb advised that developers should “measure before optimising” by inspecting site-packages within the container to identify heavy libraries.

“Being intentional about how dependencies are installed—specifically using CPU-only PyTorch—can reduce image size by ~75-80%, improve build times and deployment speed, and lower the infra and storage costs,” Jhamb wrote.