Skip to content

Commit 658959c

Browse files
committed
remove the unnecessary memory free
1 parent 8b4e9bf commit 658959c

1 file changed

Lines changed: 2 additions & 1 deletion

File tree

nunchaku/models/transformers/transformer_qwenimage.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -555,7 +555,8 @@ def forward(
555555
hidden_states = self.norm_out(hidden_states, temb)
556556
output = self.proj_out(hidden_states)
557557

558-
torch.cuda.empty_cache()
558+
if self.offload:
559+
torch.cuda.empty_cache()
559560

560561
if not return_dict:
561562
return (output,)

0 commit comments

Comments
 (0)