fix broken cuda tests (#19998)

Differential Revision: D107449731

Pull Request resolved: #19998

Author: Gasoonjia

extension/module/test/module_device_memory_test.cpp

@@ -146,17 +146,28 @@ TEST_F(ModuleDeviceMemoryTest, DeviceModelMethodMetaReportsCudaBuffer) {
   auto meta = module.method_meta("forward");
   ASSERT_TRUE(meta.ok());
 
-  // ModuleAddWithDevice has 1 planned buffer (48 bytes) on CUDA.
-  ASSERT_EQ(meta->num_memory_planned_buffers(), 1);
+  ASSERT_EQ(meta->num_memory_planned_buffers(), 2);
 
-  auto size = meta->memory_planned_buffer_size(0);
-  ASSERT_TRUE(size.ok());
-  EXPECT_EQ(size.get(), 48);
-
-  auto device = meta->memory_planned_buffer_device(0);
-  ASSERT_TRUE(device.ok());
-  EXPECT_EQ(device->type(), DeviceType::CUDA);
-  EXPECT_EQ(device->index(), 0);
+  {
+    auto size = meta->memory_planned_buffer_size(0);
+    ASSERT_TRUE(size.ok());
+    EXPECT_EQ(size.get(), 48);
+
+    auto device = meta->memory_planned_buffer_device(0);
+    ASSERT_TRUE(device.ok());
+    EXPECT_EQ(device->type(), DeviceType::CPU);
+    EXPECT_EQ(device->index(), 0);
+  }
+  {
+    auto size = meta->memory_planned_buffer_size(1);
+    ASSERT_TRUE(size.ok());
+    EXPECT_EQ(size.get(), 48);
+
+    auto device = meta->memory_planned_buffer_device(1);
+    ASSERT_TRUE(device.ok());
+    EXPECT_EQ(device->type(), DeviceType::CUDA);
+    EXPECT_EQ(device->index(), 0);
+  }
 }
 
 TEST_F(ModuleDeviceMemoryTest, DeviceModelWithSharedArenasReturnsNotSupported) {

runtime/executor/test/tensor_parser_device_test.cpp

@@ -198,8 +198,8 @@ TEST_F(TensorParserDeviceTest, CUDADeviceParsedFromPteFile) {
 
   EXPECT_EQ(cuda_tensor_count, 3)
       << "Expected 3 CUDA tensors (2 delegate inputs + 1 delegate output)";
-  EXPECT_EQ(cpu_tensor_count, 0)
-      << "Expected 0 CPU tensors (all annotated as CUDA)";
+  EXPECT_EQ(cpu_tensor_count, 3)
+      << "Expected 3 CPU tensors (2 method inputs + 1 method output)";
 }
 
 TEST_F(TensorParserDeviceTest, NonDelegatedTensorsDefaultToCPU) {
@@ -260,7 +260,7 @@ TEST_F(TensorParserDeviceTest, CudaTensorDataPtrPointsToDeviceMemory) {
   //   non_const_buffer_sizes: [0, 48]  (index 0 reserved, buffer 0 = 48 bytes)
   //   non_const_buffer_device: [{buffer_idx=1, device_type=CUDA}]
   const size_t num_buffers = method_meta->num_memory_planned_buffers();
-  ASSERT_EQ(num_buffers, 1);
+  ASSERT_EQ(num_buffers, 2);
 
   // Set up device-aware planned memory.
   std::vector<Span<uint8_t>> planned_spans;