A simple demo to train mnist in pytorch and speed up inference by TensorRT. The training code comes from here . The code to use TensorRT comes from samples in installation package of TensorRT.
The TensorRT version we use is 5.1.5. The code may not compatible with other versions of TensorRT.
You can go to （原）pytorch中使用TensorRT for more information.