Fix donated_buffer issue

[FindHao]

The previous PR #104 causes the following issue.

% python run.py --op geglu --mode fwd  --precision fp32 --metrics latency,speedup --csv --cudagraph
  0%|                                                                                                                           | 0/4 [00:03<?, ?it/s]
Caught exception, terminating early with partial results
Traceback (most recent call last):
  File "/scratch/yhao/pta/tritonbench/tritonbench/utils/triton_op.py", line 782, in run
    y_vals: Dict[str, BenchmarkOperatorMetrics] = functools.reduce(
  File "/scratch/yhao/pta/tritonbench/tritonbench/utils/triton_op.py", line 770, in _reduce_benchmarks
    acc[bm_name] = self._do_bench(
  File "/scratch/yhao/pta/tritonbench/tritonbench/utils/triton_op.py", line 981, in _do_bench
    fn = self._get_bm_func(fn_name)
  File "/scratch/yhao/pta/tritonbench/tritonbench/utils/triton_op.py", line 667, in _get_bm_func
    fwd_fn = fwd_fn_lambda(*self.example_inputs)
  File "/scratch/yhao/pta/tritonbench/tritonbench/utils/triton_op.py", line 481, in _inner
    return function(self, *args, **kwargs)
  File "/scratch/yhao/pta/tritonbench/tritonbench/operators/geglu/operator.py", line 69, in inductor_geglu
    compiled = torch.compile(self.baseline_model)
UnboundLocalError: local variable 'torch' referenced before assignment
(B, T, H)

we should use from torch._functorch import config rather than import torch._functorch.config

[facebook-github-bot]

@FindHao has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@FindHao merged this pull request in 3adf655.