Add option to profile code modules and examples

[MortenHolmRep]

Describe the solution you'd like Profiling the code allows for the potential speedup of a given module or script. Further allows for a more precise configuration of cluster submission and resource allocation when running large jobs/processes.

Tensorboard already contains a profiling tool and is installed on a new installation. An example of a profiling solution for Pytorch using Tensorboard can be found here

Alternatively Pytorch and wandb can also work together (here) to give more indepth wandb reports

[asogaard]

I think this would be really beneficial, in particular because we seem to have large overheads when running scripts, especially with W&B, and large memory usage, specifically with ParquetDataset.