Fix flash_attention test error

Differential Revision: D81511241

Pull Request resolved: #382

Author: xuzhao9

test/test_gpu/main.py

@@ -87,10 +87,10 @@ def _run_one_operator(args: List[str]):
         return
     if op.has_bwd():
         del op
-        tb_args.mode = "bwd"
         if tb_args.op in BWD_ARGS_OPS:
             args.extend(BWD_ARGS_OPS[tb_args.op])
             tb_args, extra_args = parser.parse_known_args(args)
+        tb_args.mode = "bwd"
         op = Operator(tb_args=tb_args, extra_args=extra_args)
         op.run()
         check_ci_output(op)