mirror of
https://github.com/kohya-ss/sd-scripts.git
synced 2026-04-06 13:47:06 +00:00
Change single blocks to use num_blocks value
This commit is contained in:
@@ -969,7 +969,7 @@ class Flux(nn.Module):
|
||||
def enable_block_swap(self, num_blocks: int, device: torch.device):
|
||||
self.blocks_to_swap = num_blocks
|
||||
double_blocks_to_swap = min(self.num_double_blocks - 2, num_blocks // 2)
|
||||
single_blocks_to_swap = (num_blocks - double_blocks_to_swap) * 2
|
||||
single_blocks_to_swap = (num_blocks - (num_blocks // 2)) * 2
|
||||
|
||||
assert double_blocks_to_swap <= self.num_double_blocks - 2 and single_blocks_to_swap <= self.num_single_blocks - 2, (
|
||||
f"Cannot swap more than {self.num_double_blocks - 2} double blocks and {self.num_single_blocks - 2} single blocks. "
|
||||
|
||||
Reference in New Issue
Block a user