Adding torch accelerator to FSDP2 example

Signed-off-by: dggaytan <diana.gaytan.munoz@intel.com>

Author: dggaytan

distributed/FSDP2/example.py

@@ -41,15 +41,15 @@ def main(args):
     rank = int(os.environ["LOCAL_RANK"])
     if torch.accelerator.is_available():
         device_type = torch.accelerator.current_accelerator()
-        device: torch.device = torch.device(f"{device_type}:{rank}")
+        device = torch.device(f"{device_type}:{rank}")
         torch.accelerator.device_index(rank)
         print(f"Running on rank {rank} on device {device}")
-        backend = torch.distributed.get_default_backend_for_device(device)
-        torch.distributed.init_process_group(backend=backend, device_id=device)
     else:
         device = torch.device("cpu")
         print(f"Running on device {device}")
-        torch.distributed.init_process_group(backend="gloo", device_id=device)
+
+    backend = torch.distributed.get_default_backend_for_device(device)
+    torch.distributed.init_process_group(backend=backend, device_id=device)
 
     torch.manual_seed(0)
     vocab_size = 1024