[Mod] multigpu.py debug print added

pytorch · Mar 14, 2024 · 1e81da3 · 1e81da3
1 parent a4ad613
commit 1e81da3
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/distributed/ddp-tutorial-series/multigpu.py b/distributed/ddp-tutorial-series/multigpu.py
@@ -2,6 +2,7 @@
 import torch.nn.functional as F
 from torch.utils.data import Dataset, DataLoader
 from datautils import MyTrainDataset
+from icecream import ic
 
 # --- Additional modules required for Distributed Training
 import torch.multiprocessing as mp
@@ -86,6 +87,8 @@ def prepare_dataloader(dataset: Dataset, batch_size: int):
 
 
 def main(rank: int, world_size: int, save_every: int, total_epochs: int, batch_size: int):
+    ic(rank, world_size)
+
     ddp_setup(rank, world_size)
     dataset, model, optimizer = load_train_objs()
     train_data = prepare_dataloader(dataset, batch_size)