modify test

Signed-off-by: yechank <[email protected]>

Author: yechank-nvidia

tests/unittest/_torch/multimodal/test_share_multiparams.py

@@ -53,32 +53,6 @@ def test_to_handle_unsupported_element(self):
 
         self.assertIn("Unsupported element 'multimodal_input'",
                       str(context.exception))
-        self.assertIn("Supported elements: 'multimodal_data'",
-                      str(context.exception))
-
-    def test_to_tensor_unsupported_element(self):
-        """Test to_tensor raises ValueError for unsupported elements."""
-        params = MultimodalParams()
-
-        with self.assertRaises(ValueError) as context:
-            params.to_tensor("multimodal_input")
-
-        self.assertIn("Unsupported element 'multimodal_input'",
-                      str(context.exception))
-        self.assertIn("Supported elements: 'multimodal_data'",
-                      str(context.exception))
-
-    def test_to_device_unsupported_element(self):
-        """Test to_device raises ValueError for unsupported elements."""
-        params = MultimodalParams()
-
-        with self.assertRaises(ValueError) as context:
-            params.to_device("multimodal_input", device="cuda", pin_memory=True)
-
-        self.assertIn("Unsupported element 'multimodal_input'",
-                      str(context.exception))
-        self.assertIn("Supported elements: 'multimodal_data'",
-                      str(context.exception))
 
     def test_to_tensor_basic_handle(self):
         """Test converting a basic handle back to tensor."""
@@ -122,5 +96,56 @@ def test_to_tensor_all_handles(self):
                          self.image["image_width"])
 
 
+class TestMultimodalParamsDeviceTransfer(unittest.TestCase):
+    """Test cases for to_device method in MultimodalParams."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.mm_embedding = torch.randn(3, 4, 5)
+        self.mrope_config = {
+            "mrope_rotary_cos_sin": torch.randn(2, 3),
+            "mrope_position_deltas": torch.randn(5),
+        }
+        self.image = {
+            "pixel_values": torch.randn(1, 3, 224, 224),
+            "image_height": [224],
+            "image_width": [224],
+        }
+        self.sample_multimodal_data = {
+            "multimodal_embedding": self.mm_embedding,
+            "mrope_config": self.mrope_config,
+            "image": self.image,
+        }
+
+    def test_to_device_basic(self):
+        """Test converting a basic data to device."""
+        params = MultimodalParams()
+        params.multimodal_data = {"multimodal_embedding": self.mm_embedding}
+
+        params.to_device("multimodal_data", device="cuda:0", pin_memory=True)
+
+        result = params.multimodal_data["multimodal_embedding"]
+        self.assertEqual(result.device, torch.device("cuda:0"))
+
+    def test_to_device_all_data(self):
+        """Test converting all data to device."""
+        params = MultimodalParams()
+        params.multimodal_data = self.sample_multimodal_data.copy()
+
+        params.to_device("multimodal_data", device="cuda:0", pin_memory=True)
+
+        result = params.multimodal_data["multimodal_embedding"]
+        self.assertEqual(result.device, torch.device("cuda:0"))
+
+        result = params.multimodal_data["mrope_config"]["mrope_rotary_cos_sin"]
+        self.assertEqual(result.device, torch.device("cuda:0"))
+
+        result = params.multimodal_data["mrope_config"]["mrope_position_deltas"]
+        self.assertEqual(result.device, torch.device("cuda:0"))
+
+        result = params.multimodal_data["image"]["pixel_values"]
+        self.assertEqual(result.device, torch.device("cuda:0"))
+
+
 if __name__ == "__main__":
     unittest.main()