record unit time in torch.cuda.event (#15221)

Summary: Record unit of time for torch.cuda.Event's elapsed_time

Differential Revision: D13467646

Pulled By: zou3519

fbshipit-source-id: 4f1f4ef5fa4bc5a1b4775dfcec6ab155e5bf8d6e

diff --git a/torch/cuda/streams.py b/torch/cuda/streams.py
index 16c4153..62391a5 100644 (file)
--- a/torch/cuda/streams.py
+++ b/torch/cuda/streams.py
@@ -192,7 +192,7 @@ class Event(object):
         return True
 
     def elapsed_time(self, end_event):
-        r"""Returns the time elapsed before the event was recorded."""
+        r"""Returns the time elapsed in milliseconds before the event was recorded."""
         time_ms = ctypes.c_float()
         check_error(cudart().cudaEventElapsedTime(
             ctypes.byref(time_ms), self, end_event))