Fixed misunderstanding of test_and_set

Chillee · Chillee · commit f38ea25c55b1 · 2019-05-24T16:10:25.000-07:00
diff --git a/torch/csrc/jit/passes/utils/memory_dag.cpp b/torch/csrc/jit/passes/utils/memory_dag.cpp
@@ -55,9 +55,10 @@ void collectAllContainedMemoryLocations(
     MemoryLocations& cont) {
   // we have already recursed on this element
   unsigned compIdx = Element::toIndex(elem);
-  if (cont.test_and_set(compIdx)) {
+  if (cont.test(compIdx)) {
     return;
   }
+  cont.set(compIdx);
 
   for (const auto& mem_loc : elem->getMemoryLocations()) {
     collectAllContainedMemoryLocations(Element::toElement(mem_loc), cont);
