fix a typo in the FSDP example (#1159)

pytorch · Aug 8, 2023 · 508743b · 508743b
1 parent 4440841
commit 508743b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/distributed/FSDP/utils/train_utils.py b/distributed/FSDP/utils/train_utils.py
@@ -72,7 +72,7 @@ def validation(model, rank, world_size, val_loader):
     model.eval()
     correct = 0
     local_rank = int(os.environ['LOCAL_RANK'])
-    fsdp_loss = torch.zeros(3).to(local_rank)
+    fsdp_loss = torch.zeros(2).to(local_rank)
     if rank == 0:
         inner_pbar = tqdm.tqdm(
             range(len(val_loader)), colour="green", desc="Validation Epoch"