I guessed it wouldn’t be so straightforward @tom . I have used C++ but not for a long time but I understand your warning Alex, I will not try to touch the kernel CUDA code if I am not sure about what I am doing and I don’t see myself qualified to test your solution. However, I managed to use the cast of both parameters (in python, before calling to the kernel, and after) to check the randomness reduction, as @tom suggested, but it is not yet reproducible, because during the training, the loss is quite similar, but there are peak differences, less than before, but still a little. I think that I will use this change only in the cases where the backbone has to be trained and keep it frozen otherwise.
Thanks both for your help, I have learnt a lot about the C++/CUDA PyTorch API