torchprofe broke #167
Description
The copied code from torchprof broke because pytorch profile changed. This should be implemented properly, with torchprof as a dependency (pip package) and not as a copied file from their repo which would need constant manual updating.