| 159 | @staticmethod |
| 160 | @th.cuda.amp.custom_bwd |
| 161 | def backward(ctx, *output_grads): |
| 162 | args = list(ctx.saved_tensors) |
| 163 | |
| 164 | # Filter for inputs that require grad. If none, exit early. |
| 165 | input_indices = [i for (i, x) in enumerate(args) if x.requires_grad] |
| 166 | if not input_indices: |
| 167 | return (None, None) + tuple(None for _ in args) |
| 168 | |
| 169 | with th.enable_grad(): |
| 170 | for i in input_indices: |
| 171 | if i < ctx.input_length: |
| 172 | # Not sure why the OAI code does this little |
| 173 | # dance. It might not be necessary. |
| 174 | args[i] = args[i].detach().requires_grad_() |
| 175 | args[i] = args[i].view_as(args[i]) |
| 176 | output_tensors = ctx.run_function(*args[:ctx.input_length]) |
| 177 | |
| 178 | if isinstance(output_tensors, th.Tensor): |
| 179 | output_tensors = [output_tensors] |
| 180 | |
| 181 | # Filter for outputs that require grad. If none, exit early. |
| 182 | out_and_grads = [(o, g) for (o, g) in zip(output_tensors, output_grads) if o.requires_grad] |
| 183 | if not out_and_grads: |
| 184 | return (None, None) + tuple(None for _ in args) |
| 185 | |
| 186 | # Compute gradients on the filtered tensors. |
| 187 | computed_grads = th.autograd.grad( |
| 188 | [o for (o, g) in out_and_grads], |
| 189 | [args[i] for i in input_indices], |
| 190 | [g for (o, g) in out_and_grads] |
| 191 | ) |
| 192 | |
| 193 | # Reassemble the complete gradient tuple. |
| 194 | input_grads = [None for _ in args] |
| 195 | for (i, g) in zip(input_indices, computed_grads): |
| 196 | input_grads[i] = g |
| 197 | return (None, None) + tuple(input_grads) |