From 042767613792ab467951017d8762e272812e3142 Mon Sep 17 00:00:00 2001
From: dusty-nv <dustinf@nvidia.com>
Date: Thu, 20 Sep 2018 16:55:05 +0000
Subject: [PATCH] use setDLACore()

---
 tensorNet.cpp | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/tensorNet.cpp b/tensorNet.cpp
index 5e7646e74..f5ac8a80e 100644
--- a/tensorNet.cpp
+++ b/tensorNet.cpp
@@ -111,8 +111,8 @@ static inline nvinfer1::DeviceType deviceTypeToTRT( deviceType type )
 	{
 		case DEVICE_GPU:	return nvinfer1::DeviceType::kGPU;
 		//case DEVICE_DLA:	return nvinfer1::DeviceType::kDLA;
-		case DEVICE_DLA_0:	return nvinfer1::DeviceType::kDLA0;
-		case DEVICE_DLA_1:	return nvinfer1::DeviceType::kDLA1;
+		case DEVICE_DLA_0:	return nvinfer1::DeviceType::kDLA/*0*/;
+		case DEVICE_DLA_1:	return nvinfer1::DeviceType::kDLA/*1*/;
 	}
 }
 #endif
@@ -370,6 +370,11 @@ bool tensorNet::ProfileModel(const std::string& deployFile,			   // name for caf
 
 	if( allowGPUFallback )
 		builder->allowGPUFallback(true);
+	
+	if( device == DEVICE_DLA_0 )
+		builder->setDLACore(0);
+	else if( device == DEVICE_DLA_1 )
+		builder->setDLACore(1);
 #endif
 
 	// build CUDA engine
@@ -539,7 +544,15 @@ bool tensorNet::LoadNetwork( const char* prototxt_path, const char* model_path,
 		printf(LOG_GIE "failed to create InferRuntime\n");
 		return 0;
 	}
-	
+
+#if NV_TENSORRT_MAJOR >= 5
+	// if using DLA, set the desired core before deserialization occurs
+	if( device == DEVICE_DLA_0 )
+		infer->setDLACore(0);
+	else if( device == DEVICE_DLA_1 )
+		infer->setDLACore(1);
+#endif
+
 #if NV_TENSORRT_MAJOR > 1
 	// support for stringstream deserialization was deprecated in TensorRT v2
 	// instead, read the stringstream into a memory buffer and pass that to TRT.