Skip to content

Commit d04db8d

Browse files
authored
init gds only if need (#31)
Signed-off-by: yuanyuxing.yyx <[email protected]>
1 parent 7dfe757 commit d04db8d

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

fastsafetensors/loader.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,10 @@ def __init__(
6060
global loaded_nvidia
6161
if not loaded_nvidia:
6262
fstcpp.load_nvidia_functions()
63-
if fstcpp.init_gds() != 0:
64-
raise Exception(f"[FAIL] init_gds()")
63+
if not nogds:
64+
# no need to init gds and consume 10s+ in none-gds case
65+
if fstcpp.init_gds() != 0:
66+
raise Exception(f"[FAIL] init_gds()")
6567
loaded_nvidia = True
6668
global gl_set_numa
6769
if not gl_set_numa and set_numa:

0 commit comments

Comments
 (0)