Remove redundant type conversion (#4513)

2025-03-17 05:57:35 -07:00
parent 5f9b2c62ff
commit 82dec1f70b
6 changed files with 16 additions and 10 deletions
--- a/test/srt/test_update_weights_from_tensor.py
+++ b/test/srt/test_update_weights_from_tensor.py
@@ -44,6 +44,9 @@ class TestUpdateWeightsFromTensor(unittest.TestCase):
    def test_update_weights_from_tensor(self):
        tp_sizes = [1, 2]
        for tp_size in tp_sizes:
+            if torch.cuda.device_count() < tp_size:
+                continue
+
            with self.subTest(tp_size=tp_size):
                test_update_weights_from_tensor(tp_size)