Update attention_processor.py

Author: ManfeiBai

src/diffusers/models/attention_processor.py

@@ -26,6 +26,7 @@
 
 if is_torch_xla_available():
     from torch_xla.experimental.custom_kernel import flash_attention
+    import torch_xla.distributed.spmd as xs
     import torch_xla.runtime as xr
     XLA_AVAILABLE = True
 else: