Thank you very much, I’m sure permute will come in handy in the future 
For now I fixed it by setting batch_size to None and dadding that dimension “manually”.
After another few issues, my model is training now. I hope it’s also learning