explicitly del state (#784)

This commit is contained in:
Xinyu Yang 2023-08-18 03:56:04 +08:00 committed by GitHub
parent d1744376ae
commit 73b3de79ea
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -76,6 +76,8 @@ def hf_model_weights_iterator(
state = torch.load(bin_file, map_location="cpu")
for name, param in state.items():
yield name, param
del state
torch.cuda.empty_cache()
def load_tensor_parallel_weights(