()
| 244 | out_grad_maps = {get_handle_id(t): get_grad_slot(t) for t in rlist} |
| 245 | |
| 246 | def enter_trace(): |
| 247 | for k, v in out_grad_maps.items(): |
| 248 | out_grad_maps[k] = ( |
| 249 | get_handle_id(v.grad) if v is not None else -1 |
| 250 | ) |
| 251 | backward_trace_obj.setup_env() |
| 252 | backward_trace_obj._trace.enter() |
| 253 | |
| 254 | insert_callback_as_grad_tape(enter_trace) |
| 255 | return ret |