# PyTorch ## About this Install Guide | | | |---------------------------|------------------| | Reading time: | 15 min | | Last updated: | 22 Jun 2026 | | Ecosystem dashboard: | [View](https://developer.arm.com/ecosystem-dashboard/linux?package=pytorch) | | Author: | Jason Andrews, Arm [GitHub](https://github.com/jasonrandrews) [LinkedIn](https://linkedin.com/in/jason-andrews-7b05a8) | | Official docs: | [View](https://pytorch.org/docs/stable/index.html) | | Tags: | PyTorch | This guide shows you how to install and use the tool with the most common configuration. For advanced options and complete reference information, see the official documentation. Some install guides also include optional next steps to help you explore related workflows or integrations. [PyTorch](https://pytorch.org/) is a popular end-to-end machine learning framework for Python. It is used to build and deploy neural networks, especially around tasks such as computer vision and natural language processing (NLP). Follow the instructions below to install and use PyTorch on Arm Linux. > Anaconda provides another way to install PyTorch. See the [Anaconda install guide](https://learn.arm.com/install-guides/anaconda/) to find out how to use PyTorch from Anaconda. The Anaconda version of PyTorch might be older than the version available using `pip`. ## What do I need before installing PyTorch? Confirm that you are using an Arm Linux system by running: ```bash uname -m ``` The output should be: ```bash __output__ aarch64 ``` If you see a different result, then you are not using an Arm computer running 64-bit Linux. PyTorch requires Python 3, and this can be installed with `pip`. For Ubuntu, run: ```bash sudo apt update sudo apt install python-is-python3 python3-pip python3-venv -y ``` For Amazon Linux, run: ```bash sudo dnf install python-pip -y alias python=python3 ``` ## How do I download and install PyTorch? It is recommended that you install PyTorch in your own Python virtual environment. Set up your virtual environment: ```bash python -m venv venv source venv/bin/activate ``` In your active virtual environment, install PyTorch: ```bash pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu ``` ## How do I get started with PyTorch? Test PyTorch: Use a text editor to copy and paste the code below into a text file named `pytorch.py`: ```python import torch print(torch.__version__) x = torch.rand(5, 3) print(x) exit() ``` Run the example code: ```bash python pytorch.py ``` The expected output is similar to: ```bash __output__ 2.5.1 __output__ tensor([[0.1334, 0.7932, 0.4396], [0.9409, 0.6977, 0.5904], [0.6951, 0.8543, 0.0748], [0.0293, 0.7626, 0.8668], [0.8832, 0.5077, 0.6830]]) ``` To get more information about the build options for PyTorch, run: ```bash python -c "import torch; print(*torch.__config__.show().split('\n'), sep='\n')" ``` The output will be similar to: ```bash __output__ PyTorch built with: __output__ - GCC 13.3 __output__ - C++ Version: 201703 __output__ - Intel(R) MKL-DNN v3.10.2 (Git Hash f1d471933dc852f956fd05389f9313c7148783d5) __output__ - OpenMP 201511 (a.k.a. OpenMP 4.5) __output__ - LAPACK is enabled (usually provided by MKL) __output__ - NNPACK is enabled __output__ - CPU capability usage: SVE256 ... ``` The configuration output is an advanced option to check the tools and structure used to build PyTorch. ## BFloat16 floating-point number format Recent Arm processors support the BFloat16 (BF16) number format in PyTorch. For example, AWS Graviton3 processors support BFloat16. To check if your system includes BFloat16, use the `lscpu` command: ```bash lscpu | grep bf16 ``` If the `Flags` are printed, you have a processor with BFloat16. ```bash __output__ Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm ssbs paca pacg dcpodp svei8mm svebf16 i8mm bf16 dgh rng ``` If the result is blank, you do not have a processor with BFloat16. BFloat16 provides improved performance and smaller memory footprint with the same dynamic range. You might experience a drop in model inference accuracy with BFloat16, but the impact is acceptable for the majority of applications. You can use an environment variable to enable BFloat16: ```bash export DNNL_DEFAULT_FPMATH_MODE=BF16 ``` ## LRU cache capacity LRU cache capacity is used to avoid redundant primitive creation latency overhead. This caching feature increases memory usage. If needed, you can lower the value to reduce memory usage. You should tune the capacity to an optimal value for your use case. Use an environment variable to set the value. The recommended starting value is: ```bash export LRU_CACHE_CAPACITY=1024 ``` ## Transparent huge pages Transparent huge pages (THP) provide an alternative method of utilizing huge pages for virtual memory. Enabling THP might result in improved performance because it reduces the overhead of Translation Lookaside Buffer (TLB) lookups by using a larger virtual memory page size. To check if THP is available on your system, run: ```bash cat /sys/kernel/mm/transparent_hugepage/enabled ``` The setting in brackets is your current setting. The most common output, `madvise`, is shown below: ```bash __output__ always [madvise] never ``` If the setting is `never`, you can change to `madvise` by running: ```bash echo madvise | sudo tee /sys/kernel/mm/transparent_hugepage/enabled ``` With `madvise` you can use an environment variable to check performance with and without THP. To enable THP for PyTorch: ```bash export THP_MEM_ALLOC_ENABLE=1 ``` ## Profiling example To profile a [Vision Transformer (ViT) model](https://huggingface.co/google/vit-base-patch16-224), first download the transformers and datasets libraries: ```bash pip install transformers datasets ``` Use a text editor to save the code below as `profile-vit.py`: ```python import torch from transformers import ViTFeatureExtractor, ViTForImageClassification from datasets import load_dataset from torch.profiler import profile, record_function, ProfilerActivity # Load the feature extractor and the model model_name = 'google/vit-base-patch16-224' feature_extractor = ViTFeatureExtractor.from_pretrained(model_name) model = ViTForImageClassification.from_pretrained(model_name) # Load an example image dataset = load_dataset("huggingface/cats-image", trust_remote_code=True) image = dataset["test"]["image"][0] # Preprocess the image inputs = feature_extractor(images=image, return_tensors="pt") # Perform the inference and profile it with profile(activities=[ProfilerActivity.CPU]) as prof: with record_function("mymodel_inference"): for _ in range(10): with torch.no_grad(): outputs = model(**inputs) # Print the predicted class predicted_class_idx = outputs.logits.argmax(-1).item() print(f'Predicted class: {model.config.id2label[predicted_class_idx]}') # Should be Predicted class: Egyptian cat # Print the profile print(prof.key_averages().table(sort_by="self_cpu_time_total")) ``` Run the example and check the performance information printed: ```bash python ./profile-vit.py ``` The output will be similar to: ```bash __output__ Predicted class: Egyptian cat __output__ ----------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ __output__ Name Self CPU % Self CPU CPU total % CPU total CPU time avg # of Calls __output__ ----------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ __output__ aten::addmm 72.23% 568.371ms 74.41% 585.512ms 802.072us 730 ... self CPU time total: 786.880ms ``` Experiment with the two environment variables for BFloat16 and THP and observe the performance differences. You can set each variable and run the test again and observe the new profile data and run time. ## Profiling example with dynamic quantization You can improve the performance of model inference with the `torch.nn.Linear` layer using dynamic quantization. This technique converts weights to 8-bit integers before inference and dynamically quantizes activations during inference, without the requirement for fine-tuning. However, it might impact the accuracy of your model. Use a text editor to save the code below as `profile-vit-dq.py`: ```python import torch from transformers import ViTFeatureExtractor, ViTForImageClassification from datasets import load_dataset from torch.profiler import profile, record_function, ProfilerActivity # Load the feature extractor and the model model_name = 'google/vit-base-patch16-224' feature_extractor = ViTFeatureExtractor.from_pretrained(model_name) model = ViTForImageClassification.from_pretrained(model_name) # Load an example image dataset = load_dataset("huggingface/cats-image", trust_remote_code=True) image = dataset["test"]["image"][0] # Preprocess the image inputs = feature_extractor(images=image, return_tensors="pt") # Dynamically quantize the linear layers of the model quantized_model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8) # Perform the inference and profile it with profile(activities=[ProfilerActivity.CPU]) as prof: with record_function("mymodel_inference"): for _ in range(10): with torch.no_grad(): outputs = quantized_model(**inputs) # Print the predicted class predicted_class_idx = outputs.logits.argmax(-1).item() print(f'Predicted class: {model.config.id2label[predicted_class_idx]}') # Should be Predicted class: Egyptian cat # Print the profile print(prof.key_averages().table(sort_by="self_cpu_time_total")) ``` Run the example and check the performance information printed: ```bash python ./profile-vit-dq.py ``` You should see the `quantized::linear_dynamic` layer being profiled. You can see the improvement in the model inference performance using dynamic quantization. You are now ready to use PyTorch on Arm Linux. Continue learning by exploring the many [machine learning articles and examples using PyTorch](https://pytorch.org/tutorials/).