I forgot about that aspect. Below hacky code may work for that:
class F(autograd.Function):
	
	@staticmethod
	def forward(ctx, x, g):
		with torch.enable_grad():
			x = x.clone().requires_grad_(True)
			z = g(x)
			ctx.save_for_backward(x, z)
			ctx._module = g
			return z.detach()
	
	@staticmethod
	def backward(ctx, output_grad):
		x, z = ctx.saved_tensors
		module_pars = list(ctx._module.parameters())
		g_x, *g_pars = autograd.grad(z, [x] + module_pars, output_grad)
		for p, g_p in zip(module_pars, g_pars):
			if p.grad is None:
				p.grad = g_p
			else:
				p.grad += g_p
		
		return g_x, None

Safer alternative is to extract parameters() outside and use *args, it is a bit more work to write that.