change state dict comparison to ref compare

This commit is contained in:
Jabasukuriputo Wang 2023-12-05 18:00:48 -06:00 committed by GitHub
parent f92d61497a
commit 895456c4a2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -215,7 +215,7 @@ class LoadStateDictOnMeta(ReplaceHelper):
would be on the meta device.
"""
if state_dict == sd:
if state_dict is sd:
state_dict = {k: v.to(device="meta", dtype=v.dtype) for k, v in state_dict.items()}
original(module, state_dict, strict=strict)