If I understand it correctly, you are using a single layer now and still observe the increase in loss for GPU, while loss decreases steadily on the CPU?
If this correct?
If so, could you just for the sake of debugging use torch.set_flush_denormal(True) for the CPU run and compare the results?
I don’t think this might be an issue, but don’t know any other difference between CPU and GPU handling.
If that still creates the issue, could you upload a checkpoint and input (if possible a single batch) so that we could debug it?
PS: Are you using the same data type (e.g. FP32) in both runs?
If so, could you try to run it with FP64?