mock torch_npu current_stram

ditorch/torch_npu_adapter.py

@@ -1,3 +1,17 @@
 # Copyright (c) 2024, DeepLink.
+import torch
 import torch_npu  # noqa: F401
 from torch_npu.contrib import transfer_to_npu  # noqa: F401
+
+
+def current_stream(device=None):
+    old_device = torch.cuda.current_device()
+    if device is None:
+        device = old_device
+    torch.cuda.set_device(device)
+    stream = torch_npu.npu.current_stream(device)
+    torch.cuda.set_device(old_device)
+    return stream
+
+
+torch.cuda.current_stream = current_stream

op_tools/test/test_current_stream.py

@@ -0,0 +1,18 @@
+import ditorch
+import torch
+import unittest
+
+
+class TestCurrentStream(unittest.TestCase):
+    def test_current_stream(self):
+        stream = torch.cuda.current_stream()
+        self.assertIsInstance(stream, torch.cuda.Stream)
+
+    def test_current_stream_device(self):
+        for device in range(torch.cuda.device_count()):
+            stream = torch.cuda.current_stream(device)
+            self.assertIsInstance(stream, torch.cuda.Stream)
+
+
+if __name__ == "__main__":
+    unittest.main()