Count the amount of GPU memory of intermedite variables of conv and BN during training with Pytorch

vision
1

Hi, everyone,
How to count the amount of GPU memory of intermedite variables of conv and BN during training with Pytorch? How to calculate time of single conv and BN in model during training with Pytorch? Can you give me example code? Thanks.