Removed idist.barrier() where needed (#194)

* removed idist.barrier() where needed

* added configuration for idist.barrier()

* after bash formatting

* made some more required changes

Author: sayantan1410

src/templates/template-vision-classification/data.py

@@ -12,17 +12,21 @@ def setup_data(config: Any):
     ----------
     config: needs to contain `data_path`, `train_batch_size`, `eval_batch_size`, and `num_workers`
     """
+    #::: if (it.use_dist) { :::#
     local_rank = idist.get_local_rank()
+    #::: } :::#
     transform = T.Compose(
         [
             T.ToTensor(),
             T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
         ]
     )
 
+    #::: if (it.use_dist) { :::#
     if local_rank > 0:
         # Ensure that only rank 0 download the dataset
         idist.barrier()
+    #::: } :::#
 
     dataset_train = torchvision.datasets.CIFAR10(
         root=config.data_path,
@@ -36,9 +40,12 @@ def setup_data(config: Any):
         download=True,
         transform=transform,
     )
+
+    #::: if (it.use_dist) { :::#
     if local_rank == 0:
         # Ensure that only rank 0 download the dataset
         idist.barrier()
+    #::: } :::#
 
     dataloader_train = idist.auto_dataloader(
         dataset_train,