naming nit

Author: zpcore

examples/research_projects/pytorch_xla/train_text_to_image_xla.py

@@ -520,7 +520,7 @@ def main(args):
     from torch_xla.distributed.fsdp.utils import apply_xla_patch_to_nn_linear
 
     unet = apply_xla_patch_to_nn_linear(unet, xs.xla_patched_nn_linear_forward)
-    unet.enable_use_xla_flash_attention(partition_spec=("data", None, None, None))
+    unet.enable_xla_flash_attention(partition_spec=("data", None, None, None))
 
     vae.requires_grad_(False)
     text_encoder.requires_grad_(False)

src/diffusers/models/modeling_utils.py

@@ -225,13 +225,13 @@ def fn_recursive_set_flash_attention(module: torch.nn.Module):
             if isinstance(module, torch.nn.Module):
                 fn_recursive_set_flash_attention(module)
 
-    def enable_use_xla_flash_attention(self, partition_spec: Optional[Callable] = None):
+    def enable_xla_flash_attention(self, partition_spec: Optional[Callable] = None):
         r""" 
         Enable the flash attention pallals kernel for torch_xla.
         """
         self.set_use_xla_flash_attention(True, partition_spec)
 
-    def disable_use_xla_flash_attention(self):
+    def disable_xla_flash_attention(self):
         r""" 
         Disable the flash attention pallals kernel for torch_xla.
         """