My solution for the challenge

Forgot to mentioned, I was using PyTorch 2.0 and it was game changer for both, training and inference time.

2 Likes