(dynamics, optimizer, samples)
| 38 | |
| 39 | |
| 40 | def step(dynamics, optimizer, samples): |
| 41 | loss, grads, samples, _ = l2hmc.loss_and_grads( |
| 42 | dynamics, samples, loss_fn=l2hmc.compute_loss) |
| 43 | optimizer.apply_gradients(zip(grads, dynamics.variables)) |
| 44 | |
| 45 | return loss, samples |
| 46 | |
| 47 | |
| 48 | # To be defunnable, the function cannot return an Operation, so the above |
no test coverage detected