Specify cu128 Pytorch installation (#530)

Co-authored-by: Edenzzzz <wtan45@wisc.edu>
Co-authored-by: Wenxuan Tan <wenxuan.tan@wisc.edu>
This commit is contained in:
Kevin Lin
2025-06-18 20:02:50 -05:00
committed by GitHub
co-authored by Edenzzzz Wenxuan Tan
parent 1c41a94b62
commit 78e0c7eec9
2 changed files with 5 additions and 2 deletions
+2 -2
View File
@@ -58,14 +58,14 @@ RUN source $HOME/.local/bin/env && \
# Install STA (Sliding Tile Attention)
RUN source $HOME/.local/bin/env && \
source /opt/venv/bin/activate && \
cd csrc/attn/ && \
cd csrc/attn && \
git submodule update --init --recursive && \
python setup_sta.py install
# Install VSA
RUN source $HOME/.local/bin/env && \
source /opt/venv/bin/activate && \
cd csrc/attn/ && \
cd csrc/attn && \
git submodule update --init --recursive && \
python setup_vsa.py install