Could you post a minimal code snippet using random input shapes, which would reproduce this issue, please? The error points to a missing key, but since the code seems to apparently work on the CPU, we would like to debug it on the GPU.