(*inputs, y, grad)
| 52 | raise RuntimeError("Grad shape does not match output shape") |
| 53 | |
| 54 | def f(*inputs, y, grad): |
| 55 | # Set .grad to None to avoid extra operation of gradient accumulation |
| 56 | for x in inputs: |
| 57 | if isinstance(x, torch.Tensor): |
| 58 | x.grad = None |
| 59 | y.backward(grad, retain_graph=True) |
| 60 | |
| 61 | t = benchmark.Timer( |
| 62 | stmt="f(*inputs, y=y, grad=grad)", |