Update on "Add utils to replace torchtune SDPA with ET Custom SDPA"

Differential Revision: [D67878161](https://our.internmc.facebook.com/intern/diff/D67878161)

[ghstack-poisoned]

Author: tarun292

examples/models/llama/source_transformation/sdpa.py

@@ -9,7 +9,7 @@
 # Example script for exporting Llama2 to flatbuffer
 
 import math
-from typing import Tuple, Union, Optional
+from typing import Optional, Tuple, Union
 
 import torch
 
@@ -24,7 +24,7 @@ def __init__(
         self,
         kv_cache: Optional[Union[KVCache, QuantizedKVCache]] = None,
         dim: int = -1,
-        is_causal = True,
+        is_causal=True,
     ):
         super().__init__()
         # Custom op only supports float32 currently. Converting to/from float32 is
@@ -48,8 +48,8 @@ def forward(
         k: torch.Tensor,
         v: torch.Tensor,
         bsz,
-        seqlen = None,
-        mask = None,
+        seqlen=None,
+        mask=None,
     ):
         # Custom op only supports float32 currently. Converting to/from float32 is
         # faster than not having the op.