Update spaces/zero/torch.py

spaces/zero/torch.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 
 import multiprocessing
 import os
+import time
 from concurrent.futures import ProcessPoolExecutor
 from contextlib import suppress
 from functools import partial
@@ -241,8 +242,12 @@ if (torch := maybe_import_torch()):
         bitsandbytes.unpatch()
 
     def _move(nvidia_uuid: str):
+        t0 = time.perf_counter()
         os.environ['CUDA_VISIBLE_DEVICES'] = nvidia_uuid
         torch.Tensor([0]).cuda() # CUDA init
+        t1 = time.perf_counter()
+        print("CUDA init", t1 - t0)
+        t0 = t1
         for op in to_ops.items():
             tensor, parsed_args = op
             _, dtype, _, memory_format = parsed_args
@@ -251,8 +256,17 @@ if (torch := maybe_import_torch()):
                 dtype=dtype,
                 memory_format=memory_format,
             ) # type: ignore
+        t1 = time.perf_counter()
+        print("CUDA move", t1 - t0)
+        t0 = t1
         bitsandbytes.move()
+        t1 = time.perf_counter()
+        print("BNB move", t1 - t0)
+        t0 = t1
         torch.cuda.synchronize()
+        t1 = time.perf_counter()
+        print("CUDA synchronize", t1 - t0)
+        t0 = t1
 
     def _is_in_bad_fork():
         with ProcessPoolExecutor(mp_context=multiprocessing.get_context('fork')) as e: