minor fixes

Author: zewenli98

core/runtime/TRTEngine.cpp

@@ -42,7 +42,7 @@ void* DynamicOutputAllocator::reallocateOutputAsync(
   std::vector<int64_t> shape = {static_cast<int64_t>(size)};
   auto it = buffers.find(tensorName);
   if (it == buffers.end() || it->second.sizes() != shape) {
-    buffers[tensorName] = at::empty(shape, at::TensorOptions().dtype(dtypes.at(tensorName)).device(c10::kCUDA));
+    buffers[tensorName] = at::empty(shape, at::TensorOptions().dtype(dtypes.at(tensorName)).device(at::kCUDA));
     return buffers[tensorName].data_ptr();
   } else {
     return it->second.data_ptr();

core/runtime/execute_engine.cpp

@@ -379,7 +379,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
     if (inputs.size() > 0) {
       current_device_id = inputs[0].device().index(); // Done this way to avoid a call to cudart
     } else {
-      current_device_id = c10::cuda::current_device();
+      current_device_id = at::cuda::current_device();
     }
 
     compiled_engine->caller_stream = c10::cuda::getCurrentCUDAStream(current_device_id);
@@ -428,11 +428,13 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
       for (int i = 0; i < dims.nbDims; ++i) {
         prod *= dims.d[i];
       }
-      std::vector<int64_t> dims_vec(dims.nbDims);
+      std::vector<int64_t> shape(dims.nbDims);
       for (int i = 0; i < dims.nbDims; ++i) {
-        dims_vec[i] = dims.d[i];
+        shape[i] = dims.d[i];
       }
-      output = output.reshape(-1).view(dtype).slice(0, 0, prod).reshape(dims_vec);
+      // When using the OutputAllocator, the allocated buffer might be larger than the size of the output,
+      // so we need to reshape the buffer to the output shape
+      output = output.reshape(-1).view(dtype).slice(0, 0, prod).reshape(shape);
       outputs.push_back(output);
     }
 

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

@@ -624,6 +624,8 @@ def run_output_allocator() -> torch.Tensor | Tuple[torch.Tensor, ...]:
                         .detach()
                     )
                     prod = int(torch.prod(torch.tensor(shape)))
+                    # When using the OutputAllocator, the allocated buffer might be larger than the size of the output,
+                    # so we need to reshape the buffer to the output shape
                     output = output.reshape(-1).view(dtype)[:prod].reshape(shape)
                     outputs.append(output)