update local rank value

leonnewton · Sep 16, 2023 · aa87312 · aa87312
1 parent 63be6d2
commit aa87312
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/text2vec/bertmatching_model.py b/text2vec/bertmatching_model.py
@@ -209,7 +209,7 @@ def train(
         if data_parallel:
             self.bert = nn.DataParallel(self.bert)
             num_devices = torch.cuda.device_count()
-            local_rank = int(os.environ["LOCAL_RANK"])
+            local_rank = int(os.environ.get("LOCAL_RANK", 0))
             sampler = DistributedSampler(train_dataset, num_replicas=num_devices, rank=local_rank)
             train_dataloader = DataLoader(train_dataset, batch_size=batch_size, sampler=sampler, shuffle=False)
         else:

diff --git a/text2vec/bge_model.py b/text2vec/bge_model.py
@@ -202,7 +202,7 @@ def train(
         if data_parallel:
             self.bert = nn.DataParallel(self.bert)
             num_devices = torch.cuda.device_count()
-            local_rank = int(os.environ["LOCAL_RANK"])
+            local_rank = int(os.environ.get("LOCAL_RANK", 0))
             sampler = DistributedSampler(train_dataset, num_replicas=num_devices, rank=local_rank)
             train_dataloader = DataLoader(train_dataset, batch_size=batch_size, sampler=sampler, shuffle=False)
         else:

diff --git a/text2vec/cosent_model.py b/text2vec/cosent_model.py
@@ -188,7 +188,7 @@ def train(
         if data_parallel:
             self.bert = nn.DataParallel(self.bert)
             num_devices = torch.cuda.device_count()
-            local_rank = int(os.environ["LOCAL_RANK"])
+            local_rank = int(os.environ.get("LOCAL_RANK", 0))
             sampler = DistributedSampler(train_dataset, num_replicas=num_devices, rank=local_rank)
             train_dataloader = DataLoader(train_dataset, batch_size=batch_size, sampler=sampler, shuffle=False)
         else:

diff --git a/text2vec/sentencebert_model.py b/text2vec/sentencebert_model.py
@@ -197,7 +197,7 @@ def train(
         if data_parallel:
             self.bert = nn.DataParallel(self.bert)
             num_devices = torch.cuda.device_count()
-            local_rank = int(os.environ["LOCAL_RANK"])
+            local_rank = int(os.environ.get("LOCAL_RANK", 0))
             sampler = DistributedSampler(train_dataset, num_replicas=num_devices, rank=local_rank)
             train_dataloader = DataLoader(train_dataset, batch_size=batch_size, sampler=sampler, shuffle=False)
         else: