What is the best way to implement this loss regularization term in Pytorch C++?

I am working on a loss regularization term which looks like this: image and by taking advantage of this identity
image , I implement the loss as follows(y is my noise, w is my latent, fake image is my g(w)).
std::vectortorch::Tensor g_path_regularize(const torch::Tensor& fake_img, const torch::Tensor&
latents, const torch::Tensor& mean_path_length, Scalar decay=0.01){
std::vectortorch::Tensor res;
auto noise = torch::randn_like(fake_img) / sqrt(
fake_img.size(2) * fake_img.size(3)

auto grad = torch::autograd::grad(
{(fake_img * noise).sum()}, {latents}, {}, true, true)[0];

auto path_lengths = torch::sqrt(torch::mean(torch::sum(torch::mul(grad, grad), 2),1));

auto path_mean = mean_path_length + decay * (path_lengths.mean() - mean_path_length);//tensor

auto path_penalty = torch::mean(torch::mul(path_lengths - path_mean, path_lengths - path_mean) );

res={path_penalty, path_mean.detach(), path_lengths};
return res;

And in the training process, i have a backward step on my loss regularization term and that can take a relative long time. I wonder if there are better way to implement this. Thanks!