1 files changed, 12 insertions, 2 deletions
diff --git a/tensorflow/core/common_runtime/sycl/sycl_allocator.cc b/tensorflow/core/common_runtime/sycl/sycl_allocator.cc
index 175b784825..699b54f345 100644
--- a/tensorflow/core/common_runtime/sycl/sycl_allocator.cc
+++ b/tensorflow/core/common_runtime/sycl/sycl_allocator.cc
@@ -19,16 +19,26 @@ limitations under the License.
 
 namespace tensorflow {
 
-SYCLAllocator::~SYCLAllocator() { }
+SYCLAllocator::~SYCLAllocator() {}
 
 string SYCLAllocator::Name() { return "device:SYCL"; }
 
 void *SYCLAllocator::AllocateRaw(size_t alignment, size_t num_bytes) {
+  assert(device_);
   auto p = device_->allocate(num_bytes);
   return p;
 }
 
-void SYCLAllocator::DeallocateRaw(void *ptr) { device_->deallocate(ptr); }
+void SYCLAllocator::DeallocateRaw(void *ptr) {
+  if (device_) {
+    device_->deallocate(ptr);
+  }
+}
+
+void SYCLAllocator::EnterLameDuckMode() {
+  device_->deallocate_all();
+  device_ = nullptr;
+}
 
 } // namespace tensorflow