@albanD I have a follow up question to this. So, AFAIU forward and backward pass happens inside the GPU device without holding the python GIL.
But how about the weight update operation done by an optimizer? Would optimizer access grads and params of each layer in the python environment, update the weights and push them back to the GPU? or would you translate the python logic to a cuda kernel and update the weights inside the GPU itself?