Doesn't work with nn.DataParallel

Lyken17 / pytorch-OpCounter

Count the MACs / FLOPs of your PyTorch model.

MIT License

4.9k stars 528 forks source link

Doesn't work with nn.DataParallel #131

Closed danielamassiceti closed 3 years ago

danielamassiceti commented 3 years ago

Hi there - I just wanted to flag that I ran into an issue when my model was wrapped around a nn.DataParallel - thop gave an error that said the model was on a different device (cpu) to the input (gpu)

Lyken17 commented 3 years ago

If you are using dataparalle, then you can profile(m.module, blabla) to perform the similar function.