From 9189deb24164680421c6d4ce8adbe585dc0c96d3 Mon Sep 17 00:00:00 2001 From: Guangda Lai Date: Tue, 2 Jul 2019 17:36:26 -0700 Subject: [PATCH] Use CapturableResource._resource_device instead of creating a new member for the device. PiperOrigin-RevId: 256273685 --- tensorflow/python/compiler/tensorrt/trt_convert.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tensorflow/python/compiler/tensorrt/trt_convert.py b/tensorflow/python/compiler/tensorrt/trt_convert.py index 1c7455df992..d25b1477886 100644 --- a/tensorflow/python/compiler/tensorrt/trt_convert.py +++ b/tensorflow/python/compiler/tensorrt/trt_convert.py @@ -828,18 +828,16 @@ class TRTEngineResource(tracking.TrackableResource): filename, maximum_cached_engines, device="GPU"): - super( - TRTEngineResource, - self).__init__(deleter=TRTEngineResourceDeleter(resource_name, device)) + super(TRTEngineResource, self).__init__( + device=device, deleter=TRTEngineResourceDeleter(resource_name, device)) self._resource_name = resource_name # Track the serialized engine file in the SavedModel. self._filename = self._track_trackable( tracking.TrackableAsset(filename), "_serialized_trt_engine_filename") self._maximum_cached_engines = maximum_cached_engines - self._device = device def _create_resource(self): - return _get_resource_handle(self._resource_name, self._device) + return _get_resource_handle(self._resource_name, self._resource_device) def _initialize(self): gen_trt_ops.populate_trt_engine_cache(