[opt] set retry_on_failure to false for tf allocator

LinGeLin · May 20, 2024 · 34cc36d · 34cc36d
1 parent 53b5ac8
commit 34cc36d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/...flow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_hkv.h b/...flow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_hkv.h
@@ -354,9 +354,11 @@ class TFOrDefaultAllocator : public nv::merlin::BaseAllocator {
   void alloc(const NMMemType type, void** ptr, size_t size,
              unsigned int pinned_flags = cudaHostAllocDefault) override {
     if (!use_default_allocator_) {
+      tensorflow::AllocationAttributes allocation_attr(false, false, nullptr);
       switch (type) {
         case NMMemType::Device:
-          *ptr = tf_device_allocator_->AllocateRaw(kAllocatorAlignment, size);
+          *ptr = tf_device_allocator_->AllocateRaw(kAllocatorAlignment, size,
+                                                   allocation_attr);
           if (nullptr == *ptr) {
             throw std::runtime_error(
                 "Failed to allocator gpu memory, please adjust param 'max_hbm' "