Refactor module_size function.

This commit is contained in:
comfyanonymous 2026-01-05 03:40:49 -05:00 committed by GitHub
parent d1b9822f74
commit 1a30002a41
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -456,7 +456,7 @@ def module_size(module):
sd = module.state_dict()
for k in sd:
t = sd[k]
module_mem += t.nelement() * t.element_size()
module_mem += t.nbytes
return module_mem
class LoadedModel: