Set kIntsonDeviceAttr for native function in TRTEngineOp
This is required to make sure that all TRTEngineOp outputs are always on a GPU device as expected. Otherwise, some of the tensors of type DT_INT32 would be on host if the op generating the tensor has host memory tag set. PiperOrigin-RevId: 231707027
Loading
Please sign in to comment