Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
(x, t_input)
| 278 | return -expand_dims(sigma_t, dims) * output |
| 279 | |
| 280 | def cond_grad_fn(x, t_input): |
| 281 | """ |
| 282 | Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t). |
| 283 | """ |
| 284 | with torch.enable_grad(): |
| 285 | x_in = x.detach().requires_grad_(True) |
| 286 | log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs) |
| 287 | return torch.autograd.grad(log_prob.sum(), x_in)[0] |
| 288 | |
| 289 | def model_fn(x, t_continuous): |
| 290 | """ |