| 275 | return op |
| 276 | |
| 277 | def decorated_func(*args): |
| 278 | def apply_expr(op, *args, nr_out=None): |
| 279 | op = as_op(op, len(args)) |
| 280 | results = apply(op, *args) |
| 281 | if nr_out is None: |
| 282 | assert len(results) == 1 |
| 283 | return results[0] |
| 284 | else: |
| 285 | assert len(results) == nr_out |
| 286 | return results |
| 287 | |
| 288 | def apply_const(value, dtype=dtype, device=device): |
| 289 | return Const(value, dtype, device) |
| 290 | |
| 291 | outputs, outputs_has_grad = func(args, apply_expr, apply_const) |
| 292 | outputs = [ |
| 293 | output if has_grad else output.detach() |
| 294 | for output, has_grad in zip(outputs, outputs_has_grad) |
| 295 | ] |
| 296 | return outputs |
| 297 | |
| 298 | return decorated_func |
| 299 | |