lucidrains / vit-pytorch

Implementation of Vision Transformer, a simple way to achieve SOTA in vision classification with only a single transformer encoder, in Pytorch
MIT License
19.93k stars 2.99k forks source link

cats_and_dogs.inpynb can't use GPU training #77

Open FalconJK opened 3 years ago

FalconJK commented 3 years ago

I tried to run example/cats_and_dogs.inpynb on my V100 GPU Data and models are in the GPU memorys But the usage of GPU is always 0%, the other way, CPU's usage are in 900%~1000% I am not sure what device am i training on

p.s. I tried to print the ViT model device ViT.device, it tells me device is not define

cristianoc20 commented 2 years ago

hi, i meet the same issue.Have you solve it?