Traced model output differs on C++ and Python
Labels: oncall: jit

### 🐛 Describe the bug

Hey guys. Currently, I'm facing a problem related to the output of a model in C++ and Python.

The model receives a set of PointCloud voxels and outputs a feature map to decode the 3D bounding box. I'll provide the model to reproduce the issue.

I'm converting the PyTorch model to JIT as follows:
```
pillars = np.load("pillars.npy")
coors_batch = np.load("coors_batch.npy")
npoints_per_pillar = np.load("npoints_per_pillar.npy")

example_input = (
    pillars, coors_batch, npoints_per_pillar
)

traced_model = torch.jit.trace(model_torch, example_input)
# traced_script_module = torch.jit.script(model, example_input)
torch.jit.save(traced_model, "point_pillar.pt")
# print("[INFO] Model converted.")

model_traced = torch.jit.load("point_pillar.pt")

out1_trace, out2_trace, out3_trace = model_traced(pillars, coors_batch, npoints_per_pillar)

print('[INFO] Class Pred Shape = {}'.format(out1_trace.shape))
print('[INFO] Class Pred max = {}'.format(out1_trace.max()))
print('[INFO] Box Pred Shape = {}'.format(out2_trace.shape))
print('[INFO] Box Pred max = {}'.format(out2_trace.max()))
print('[INFO] Direction Pred Shape = {}'.format(out3_trace.shape))
print('[INFO] Direction Pred max = {}'.format(out3_trace.max()))

print("[INFO] Out1 = {}".format(torch.abs(out1_trace - out1_torch).max())) 
print("[INFO] Out2 = {}".format(torch.abs(out2_trace - out2_torch).max()))
print("[INFO] Out3 = {}".format(torch.abs(out3_trace - out3_torch).max()))

```

The above code outputs 0.0 for all the last lines, i.e. no difference at all.
```
[INFO] Class Pred Shape = torch.Size([1, 18, 248, 216])
[INFO] Class Pred max = 2.3731374740600586
[INFO] Box Pred Shape = torch.Size([1, 42, 248, 216])
[INFO] Box Pred max = 1.7485209703445435
[INFO] Direction Pred Shape = torch.Size([1, 12, 248, 216])
[INFO] Direction Pred max = 7.032490253448486

[INFO] Out1 = 0.0
[INFO] Out2 = 0.0
[INFO] Out3 = 0.0
```

However, when I load the same traced model in C++ with Libtorch, I receive a different output. Here's a sample:

```
#include <torch/script.h> // One-stop header.
#include <iostream>
#include <memory>
#include "inference/PointPillars.h"
#include <chrono>
#include <fstream>
#include <c10/cuda/CUDACachingAllocator.h>
#include <Eigen/Dense>
#include <cassert>


int main(int argc, const char* argv[]) {
    std::string modelPath = "./point_pillar.pt";

    torch::NoGradGuard no_grad_guard;

    bool CUDA = true;
    model = torch::jit::load(modelPath);
    
    detector->preprocess(xyzi, inputs);

    for (auto input : inputs) {
        std::cout << "Inputs = " << input.toTensor().sizes() << std::endl;
        std::cout << "Inputs max = " << input.toTensor().max() << std::endl;
    }

    // Allocating placeholders
    torch::Tensor clsPred;
    torch::Tensor boxPred;
    torch::Tensor dirPred;

    detector->detect(inputs, clsPred, boxPred, dirPred);

    std::cout << "[INFO] Class Pred shape = " << clsPred.sizes() << std::endl;
    std::cout << "[INFO] Class Pred max = " << clsPred.max() << std::endl;
    std::cout << "[INFO] Box Pred shape = " << boxPred.sizes() << std::endl;
    std::cout << "[INFO] Box Pred max = " << boxPred.max() << std::endl;
    std::cout << "[INFO] Direction Pred shape = " << dirPred.sizes() << std::endl;
    std::cout << "[INFO] Direction Pred max = " << dirPred.max() << std::endl;

    std::cout << "[INFO] Done inference." << std::endl;

    return 0;
}
```

That outputs the following tensors:

```
INFO] Class Pred shape = [1, 18, 248, 216]
[INFO] Class Pred max = -4.45092
[ CUDAFloatType{} ]
[INFO] Box Pred shape = [1, 42, 248, 216]
[INFO] Box Pred max = 1.36825
[ CUDAFloatType{} ]
[INFO] Direction Pred shape = [1, 12, 248, 216]
[INFO] Direction Pred max = 7.03276
[ CUDAFloatType{} ]
```

At the end, the last tensor, `Direction`, seems to be close, but the others aren't.

- Versions:
    - The Python PyTorch version that I'm using is `'1.9.0+cu111'`.
    - The Libtorch Version is `libtorch-cxx11-abi-shared-with-deps-1.13.0+cu117`.



### Versions

Collecting environment information...
PyTorch version: 1.9.0+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.23.1
Libc version: glibc-2.31

Python version: 3.8.10 (default, Jun 22 2022, 20:18:18)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-58-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: 
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3060 Laptop GPU
Nvidia driver version: 520.61.05
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.7.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.24.1
[pip3] torch==1.9.0+cu111
[pip3] torch2trt==0.4.0
[pip3] torchaudio==0.9.0
[pip3] torchvision==0.10.0+cu111
[conda] Could not collect

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel