Faster Gradient Penalty Calculation

[AhmedThahir]

This is referring to: https://github.com/t-vi/pytorch-tvmisc/blob/master/misc/gradient_regularized_gan_gaussian_toy_problem.ipynb

Current Code

grads = torch.autograd.grad(scores, discriminator.parameters(), retain_graph=True, create_graph=True)
gr_norm_sq = 0.0
for gr in grads:
    gr_norm_sq += (gr**2).sum()

Suggestion

grads_norm_penalty = (
	torch.concat(
		[g.view(-1) for g in grads]
	)
	.norm(p=2)
)

Question

What is your opinion on this? This tries to batch operations as much as possible.