lint

Author: Rachmanino

examples/flash_attention/example_mha_bwd.py

@@ -352,4 +352,4 @@ def run1():
     parser.add_argument('--d_head', type=int, default=64, help='Head dimension')
     parser.add_argument('--causal', type=bool, default=False, help='Causal flag')
     args = parser.parse_args()
-    main(args.batch, args.h, args.n_ctx, args.d_head, args.causal)
+    main(args.batch, args.h, args.n_ctx, args.d_head, args.causal)

examples/flash_attention/example_mha_bwd_wgmma_pipelined.py

@@ -362,4 +362,4 @@ def run1():
     parser.add_argument('--d_head', type=int, default=64, help='Head dimension')
     parser.add_argument('--causal', type=bool, default=False, help='Causal flag')
     args = parser.parse_args()
-    main(args.batch, args.h, args.n_ctx, args.d_head, args.causal)
+    main(args.batch, args.h, args.n_ctx, args.d_head, args.causal)