mrq/vall-e
5
9
Fork 0
Code Issues 8 Pull Requests Packages Projects Releases Wiki Activity

oops

Browse Source
This commit is contained in:
mrq 2024-08-01 22:56:04 -05:00
parent 443422ecb5
commit 808a79ebaf
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
vall_e/utils/utils.py
View File

@ -455,7 +455,7 @@ def get_model_offload_policy(module, policy=None):
# check if within budget # check if within budget
if buffered_size - size >= 0: if buffered_size - size >= 0:
# add to buffer # add to buffer
buffered_modules.append( name ) buffered_modules.append( (name, size) )
buffered_size -= size buffered_size -= size
# budget exceeded, flush buffer # budget exceeded, flush buffer
else: else:
@ -465,7 +465,7 @@ def get_model_offload_policy(module, policy=None):
break break
if buffered_modules and buffered_size: if buffered_modules and buffered_size:
device["modules"] += buffered_modules device["modules"] += [ name for name, size in buffered_modules ]
device["free"] = buffered_size device["free"] = buffered_size
modules = discarded modules = discarded