refactor: Addressing PR comments

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

docsrc/tutorials/use_from_pytorch.rst

@@ -42,6 +42,7 @@ at the documentation for the TRTorch ``TensorRTCompileSpec`` API.
                 "device": {
                     "device_type": trtorch.DeviceType.GPU,
                     "gpu_id": 0,
+                    "dla_core": 0,
                     "allow_gpu_fallback": True
                 },
                 "capability": trtorch.EngineCapability.default,

py/trtorch/_compile_spec.py

@@ -195,10 +195,10 @@ def TensorRTCompileSpec(compile_spec: Dict[str, Any]) -> torch.classes.tensorrt.
                             } # Dynamic input shape for input #2
                         ],
                         "device": {
-                        "device_type": torch.device("cuda"), # Type of device to run engine on (for DLA use trtorch.DeviceType.DLA)
-                        "gpu_id": 0, # Target gpu id to run engine (Use Xavier as gpu id for DLA)
-                        "dla_core": 0, # (DLA only) Target dla core id to run engine
-                        "allow_gpu_fallback": false, # (DLA only) Allow layers unsupported on DLA to run on GPU
+                            "device_type": torch.device("cuda"), # Type of device to run engine on (for DLA use trtorch.DeviceType.DLA)
+                            "gpu_id": 0, # Target gpu id to run engine (Use Xavier as gpu id for DLA)
+                            "dla_core": 0, # (DLA only) Target dla core id to run engine
+                            "allow_gpu_fallback": false, # (DLA only) Allow layers unsupported on DLA to run on GPU
                         },
                         "op_precision": torch.half, # Operating precision set to FP16
                         "refit": False, # enable refit

py/trtorch/csrc/tensorrt_classes.h

@@ -69,7 +69,7 @@ struct Device : torch::CustomClassHolder {
         allow_gpu_fallback(false) // allow_gpu_fallback
   {}
 
-  ADD_ENUM_GET_SET(device_type, DeviceType, 1);
+  ADD_ENUM_GET_SET(device_type, DeviceType, static_cast<int64_t>(DeviceType::kDLA));
   ADD_FIELD_GET_SET(gpu_id, int64_t);
   ADD_FIELD_GET_SET(dla_core, int64_t);
   ADD_FIELD_GET_SET(allow_gpu_fallback, bool);
@@ -98,11 +98,11 @@ struct CompileSpec : torch::CustomClassHolder {
     device = *d;
   }
 
-  ADD_ENUM_GET_SET(op_precision, DataType, 2);
+  ADD_ENUM_GET_SET(op_precision, DataType, static_cast<int64_t>(DataType::kChar));
   ADD_FIELD_GET_SET(refit, bool);
   ADD_FIELD_GET_SET(debug, bool);
   ADD_FIELD_GET_SET(strict_types, bool);
-  ADD_ENUM_GET_SET(capability, EngineCapability, 2);
+  ADD_ENUM_GET_SET(capability, EngineCapability, static_cast<int64_t>(EngineCapability::kSAFE_DLA));
   ADD_FIELD_GET_SET(num_min_timing_iters, int64_t);
   ADD_FIELD_GET_SET(num_avg_timing_iters, int64_t);
   ADD_FIELD_GET_SET(workspace_size, int64_t);

tests/py/test_to_backend_api.py

@@ -22,6 +22,7 @@ def setUp(self):
                     "device": {
                         "device_type": trtorch.DeviceType.GPU,
                         "gpu_id": 0,
+                        "dla_core": 0,
                         "allow_gpu_fallback": True
                     },
                     "capability": trtorch.EngineCapability.default,