Easily benchmark PyTorch model FLOPs, latency, throughput, max allocated memory and energy consumption in one go.
pip install pytorch-benchmark