Skip to content

Commit 094e7af

Browse files
dcasbolNicolasHug
andauthored
Consider even block sizes in drop block 2d (#9157)
Co-authored-by: Nicolas Hug <nh.nicolas.hug@gmail.com>
1 parent b208f7f commit 094e7af

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

torchvision/ops/drop_block.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,9 @@ def drop_block2d(
3636

3737
N, C, H, W = input.size()
3838
block_size = min(block_size, W, H)
39+
if block_size % 2 == 0:
40+
raise ValueError(f"block size should be odd. Got {block_size} which is even.")
41+
3942
# compute the gamma of Bernoulli distribution
4043
gamma = (p * H * W) / ((block_size**2) * ((H - block_size + 1) * (W - block_size + 1)))
4144
noise = torch.empty((N, C, H - block_size + 1, W - block_size + 1), dtype=input.dtype, device=input.device)
@@ -82,6 +85,9 @@ def drop_block3d(
8285

8386
N, C, D, H, W = input.size()
8487
block_size = min(block_size, D, H, W)
88+
if block_size % 2 == 0:
89+
raise ValueError(f"block size should be odd. Got {block_size} which is even.")
90+
8591
# compute the gamma of Bernoulli distribution
8692
gamma = (p * D * H * W) / ((block_size**3) * ((D - block_size + 1) * (H - block_size + 1) * (W - block_size + 1)))
8793
noise = torch.empty(

0 commit comments

Comments
 (0)