Skip to content

Commit 97a7f51

Browse files
lsteinLincoln Stein
andauthored
don't use cpu state_dict for model unpatching when executing on cpu (#6631)
Co-authored-by: Lincoln Stein <lstein@gmail.com>
1 parent 0583101 commit 97a7f51

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

invokeai/backend/model_manager/load/model_cache/model_cache_default.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,8 @@ def put(
167167
size = calc_model_size_by_data(self.logger, model)
168168
self.make_room(size)
169169

170-
state_dict = model.state_dict() if isinstance(model, torch.nn.Module) else None
170+
running_on_cpu = self.execution_device == torch.device("cpu")
171+
state_dict = model.state_dict() if isinstance(model, torch.nn.Module) and not running_on_cpu else None
171172
cache_record = CacheRecord(key=key, model=model, device=self.storage_device, state_dict=state_dict, size=size)
172173
self._cached_models[key] = cache_record
173174
self._cache_stack.append(key)

0 commit comments

Comments
 (0)