Recommend Nvidia's Pytorch container

f0c40b7d · Tri Dao · 3cad2ab3 · f0c40b7d
Commit f0c40b7d authored May 19, 2023 by Tri Dao
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

README.md README.md +4 -0

No files found.
--- a/README.md
+++ b/README.md
@@ -48,6 +48,10 @@ Requirements:
 - CUDA 11.4 and above.
 - PyTorch 1.12 and above.

+We recommend the
+[Pytorch](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch)
+container from Nvidia, which has all the required tools to install FlashAttention.
+
 To install:
 ```sh
 pip install flash-attn