search

pytorch-labs / attention-gym

Helpful tools and examples for working with flex-attention
BSD 3-Clause "New" or "Revised" License
434 stars 22 forks source link

torch.nn.attention folder seems to be missing on torch 2.5 #36

Closed kchl5 closed 1 month ago

kchl5 commented 1 month ago

I was trying to install flex_attention and followed the proposed installation path, also I have installed attention-gym:

I have tried both highly and on torch 2.50, see below:

torch 2.5.0.dev20240910+cu121 torchaudio 2.5.0.dev20240910+cu121 torchmetrics 1.4.1 torchvision 0.20.0.dev20240910+cu121

Even with both of these approaches, I encounter the following errors:

from torch.nn.attention.flex_attention import create_block_mask

ModuleNotFoundError: No module named 'torch.nn.attention.flex_attention'

drisspg commented 1 month ago

hmm that is strange, let me double check this

drisspg commented 1 month ago
Screenshot 2024-09-11 at 9 47 32 AM

I am unable to repro this ... hmm

kchl5 commented 1 month ago

Thanks for looking into this. I have managed to reproduce your code by only installing torch 2.5 using the command below:

pip3 install torch==2.5.0 --index-url https://down
load.pytorch.org/whl/test/cu121

I finally managed to install it. For my last attempt, I also installed torchvision and torchaudio which might have changed the torch installation. Hence, I am closing the issue.