I’m not familiar with the einops repository so don’t know how the different orders could create the issue. To further isolate it, you could rerun the code with CUDA_LAUNCH_BLOCKING=1 and see, if the error message would point to another operation, as cuDNN might be running into a previous (sticky) error.
If no error checking is done in the 3rd party libraries, you might also need to synchronize the code manually and e.g. print a CUDATensor to check, if an internal assert was triggered.