"docs/vscode:/vscode.git/clone" did not exist on "3e59b1d477d4876da53d7b3845732842cff3c66e"
Unverified Commit dc67d976 authored by Yineng Zhang's avatar Yineng Zhang Committed by GitHub
Browse files

misc: speedup load safetensors (#1319)


Co-authored-by: default avatarispobock <ISPObaoke@163.com>
parent 1e495e08
...@@ -162,6 +162,7 @@ class ModelRunner: ...@@ -162,6 +162,7 @@ class ModelRunner:
return min_per_gpu_memory return min_per_gpu_memory
def load_model(self): def load_model(self):
torch.set_num_threads(1)
logger.info( logger.info(
f"Load weight begin. avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB" f"Load weight begin. avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
) )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment