Hi,
If G = dce_loss / dW1 where ce_loss is the cross entropy loss. Then you can compute this as:
ce_loss = criterion(output, target)
G = autograd.grad(ce_loss, W1, create_graph=True)[0]
loss = ce_loss + (G * W1).pow(2).sum()