Skip to content

Commit 915fdb5

Browse files
Fix lowvram edge case.
1 parent 5a8a489 commit 915fdb5

1 file changed

Lines changed: 4 additions & 0 deletions

File tree

comfy/ops.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -264,10 +264,14 @@ def fp8_linear(self, input):
264264
scale_input = self.scale_input
265265
if scale_weight is None:
266266
scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
267+
else:
268+
scale_weight = scale_weight.to(input.device)
269+
267270
if scale_input is None:
268271
scale_input = torch.ones((), device=input.device, dtype=torch.float32)
269272
inn = input.reshape(-1, input.shape[2]).to(dtype)
270273
else:
274+
scale_input = scale_input.to(input.device)
271275
inn = (input * (1.0 / scale_input).to(input.dtype)).reshape(-1, input.shape[2]).to(dtype)
272276

273277
if bias is not None:

0 commit comments

Comments
 (0)