From 808a79ebaf7a5cd04145673e2e988157fcdcd3de Mon Sep 17 00:00:00 2001
From: mrq <mrq@ecker.tech>
Date: Thu, 1 Aug 2024 22:56:04 -0500
Subject: [PATCH] oops

---
 vall_e/utils/utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/vall_e/utils/utils.py b/vall_e/utils/utils.py
index 6227da4..f4c4d85 100755
--- a/vall_e/utils/utils.py
+++ b/vall_e/utils/utils.py
@@ -455,7 +455,7 @@ def get_model_offload_policy(module, policy=None):
 				# check if within budget
 				if buffered_size - size >= 0:
 					# add to buffer
-					buffered_modules.append( name )
+					buffered_modules.append( (name, size) )
 					buffered_size -= size
 				# budget exceeded, flush buffer
 				else:
@@ -465,7 +465,7 @@ def get_model_offload_policy(module, policy=None):
 					break
 
 			if buffered_modules and buffered_size:
-				device["modules"] += buffered_modules
+				device["modules"] += [ name for name, size in buffered_modules ]
 				device["free"] = buffered_size
 
 		modules = discarded