fixed issue with local_rank on colab

sileod · Jun 30, 2023 · aa7edca · aa7edca
1 parent a6c7489
commit aa7edca
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/src/tasknet/models.py b/src/tasknet/models.py
@@ -452,12 +452,8 @@ def get_single_train_dataloader(self, task_name, train_dataset):
         """
         if self.train_dataset is None:
             raise ValueError("Trainer: training requires a train_dataset.")
-        train_sampler = (
-            RandomSampler(train_dataset)
-            if self.args.local_rank == -1
-            else DistributedSampler(train_dataset)
-        )
-
+        train_sampler = (RandomSampler(train_dataset) if torch.cuda.device_count()<2 or self.args.local_rank == -1 else DistributedSampler(train_dataset))
+
         data_loader = DataLoaderWithTaskname(
             task_name=task_name,
             data_loader=DataLoader(