You don’t need to do anything special. What you are doing is fine but if you want to check gradients inside that function you need to rename log_sigma_A as I did.
log_sigma_A = log_sigma_AA.cuda() Here you can check log_sigma_AA.grad
log_sigma_A = log_sigma_A.cuda() Here you cannot check log_sigma_A.grad as log_sigma_A became a non-leaf tensor which point to outer log_sigma_A. It works but cannot point to original log_sigma_A