INNER CODE UNIT · Python

_set_device

onnx/onnx-tensorrt · onnx_tensorrt/backend.py:137

    def _set_device(self, device):
        self.device = device
        assert(device.type == DeviceType.CUDA)
        cudaSetDevice(device.device_id)

    def _deserialize(self, trt_blob):
        self.runtime = trt.Runtime(TRT_LOGGER)
        del self.parser # Parser no longer needed for ownership of plugins
        trt_engine = self.runtime.deserialize_cuda_engine(trt_blob)
        return trt_engine

    def run(self, inputs, **kwargs):
        """Execute the prepared engine and return the outputs as a named tuple.
        inputs -- Input tensor(s) as a Numpy array or list of Numpy arrays.
        """
        if isinstance(inputs, np.ndarray):
            inputs = [inputs]

View source record →