fix(torch): avoid race condition when building alphabet

jolibrain · Jan 25, 2023 · b1accb7 · b1accb7
1 parent 9541de1
commit b1accb7
Showing 1 changed file with 14 additions and 4 deletions.
diff --git a/src/backends/torch/torchdataset.cc b/src/backends/torch/torchdataset.cc
@@ -890,11 +890,21 @@ namespace dd
             target_tensor[i] = it->second;
           }
         else if (!_test)
+#pragma omp ordered
           {
-            this->_logger->info("added {} to alphabet", c);
-            int id = alphabet.size();
-            alphabet[c] = id;
-            target_tensor[i] = id;
+            // in a parallel loop, recheck in order to avoid race condition
+            auto it = alphabet.find(c);
+            if (it != alphabet.end())
+              {
+                target_tensor[i] = it->second;
+              }
+            else
+              {
+                this->_logger->info("added {} to alphabet", c);
+                int id = alphabet.size();
+                alphabet[c] = id;
+                target_tensor[i] = id;
+              }
           }
         else
           {