hszhao / SAN

Exploring Self-attention for Image Recognition, CVPR2020.
MIT License
747 stars 132 forks source link

can this network be deployed on cpus? #2

Closed iamZe closed 4 years ago

iamZe commented 4 years ago

thanks.

bluesky314 commented 4 years ago

+1

bluesky314 commented 4 years ago

@hszhao Any reply to this? I tried training on cpu and it did not work.

hszhao commented 4 years ago

Hi, currently only GPU version is supported as CPU training is much slower when compared to GPU training. To enable CPU training, you need to remove the CUDA part in 'tool/train.py' and modify the operators in 'model/san.py'.