Get a summary regarding memory required to train my model?(Both GPU and CPU)

Is there any summary type function that could give me an insight regarding how much memory would be required accordingly with the layers and parameters I would have set in my model.This is regarding cuda memory but if there is something like this for cpu also it would be great if someone could tell me. Thanks in advance :smile: