| 205 | del fargs |
| 206 | |
| 207 | def exit_trace(): |
| 208 | backward_trace_obj._trace.exit() |
| 209 | backward_trace_obj.unset_env() |
| 210 | for k, v in inp_grad_maps.items(): |
| 211 | inp_grad_maps[k] = ( |
| 212 | get_handle_id(v.grad) if v is not None else -1 |
| 213 | ) |
| 214 | |
| 215 | # if we want to trace the gradmanager attach callback, we will replace |
| 216 | # the original grad by the grad processed the gradmanger attach callback |
| 217 | if trace_gmcallback: |
| 218 | |
| 219 | def cb_wrapper(cb): |
| 220 | def wrapper(param, grad): |
| 221 | return ( |
| 222 | grad |
| 223 | if grad._is_external_value() |
| 224 | else cb(param, grad) |
| 225 | ) |
| 226 | |
| 227 | return wrapper |
| 228 | |
| 229 | current_gm = get_backwarding_grad_manager() |
| 230 | for _attached_tensor_id, grad in current_gm._gradients.items(): |
| 231 | spec = current_gm._attach_specs.get(_attached_tensor_id) |
| 232 | _attached_tensor = spec and spec.tensor() |
| 233 | if _attached_tensor is None: |
| 234 | continue |
| 235 | inp_hid = get_handle_id(_attached_tensor) |
| 236 | if inp_hid in inp_grad_maps: |
| 237 | spec.callbacks = list(map(cb_wrapper, spec.callbacks)) |
| 238 | inp_grad_maps[inp_hid] = get_handle_id(grad) |
| 239 | |
| 240 | insert_callback_as_grad_tape(exit_trace) |
| 241 | |