How can I deal with data race when I am creating a taichi-based hash map? The code runs successfully on CPU single thread.

[pyh3552]

@ti.func
    def find(self, k: ti.template(), create: bool = False) -> ti.int32:
        h = self._hash(k) % self.capacity_
        return_value = 0
        while True:
            e = ti.atomic_add(self.table_[h], 0)
            if (e == -1):
                if (create):
                    tmp = ti.atomic_add(self.filled_[None], 1)
                    for i in range(self.feature_dim_):
                        ti.atomic_add(self.features_[tmp * self.feature_dim_ + i], k[i])
                    self.table_[h] = tmp
                    return_value = tmp
                    # print("created")
                    break
                else: 
                    # ti.atomic_add(return_value, e)
                    return_value = ti.atomic_add(e, 0)
                    # print("not exist")
                    break
            good = True
            for i in range(self.feature_dim_):
                if good and self.features_[e * self.feature_dim_ + i] != k[i]:  
                    good = False
            if good:
                return_value = ti.atomic_add(e, 0)
                # print("found")
                break
            h += 1
            if (h == self.capacity_):
                h = 0
        # print(f"===============return{self.filled_[None]}=============")
        # print(return_value)
        return return_value