hoya012 / automatic-mixed-precision-tutorials-pytorch

Automatic Mixed Precision Tutorials using pytorch. Based on PyTorch 1.6 Official Features, implement classification codebase using custom dataset.
MIT License
88 stars 6 forks source link

Models are NOT pushed to multiple GPU's #2

Open gbrow004 opened 3 years ago

gbrow004 commented 3 years ago

I have 2x GPU's and the script identifies both of them and reports "Model pushed to 2 GPU(s)" , type geforce etc. However, only 1x GPU is being used. At the very least, I think you need to wrap the model in the nn.DataParallel wrapper and add an autocast.