Pytorch- Vectorisation of code

Hi Malcolm.
Thanks for the prompt response and also for checking back.
On just reading, it seems to be exactly what I was looking for.
Because of some personal constraints, I haven’t been able to check the correctness and timing improvement of the solution yet. I’ll get back with those when possible for me.
Thanks a lot :slight_smile: