mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Fix CUDA crash w/ channels-last + CSP models. Remove use of chunk()

Browse Source
This commit is contained in:
Ross Wightman 2021-02-23 13:15:52 -08:00
parent 317ea3e599
commit 7ef7788ee9
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
timm/models/cspnet.py
View File

@ -264,9 +264,11 @@ class CrossStage(nn.Module):
if self.conv_down is not None:
x = self.conv_down(x)
x = self.conv_exp(x)
xs, xb = x.chunk(2, dim=1)
split = x.shape[1] // 2
xs, xb = x[:, :split], x[:, split:]
xb = self.blocks(xb)
out = self.conv_transition(torch.cat([xs, self.conv_transition_b(xb)], dim=1))
xb = self.conv_transition_b(xb).contiguous()
out = self.conv_transition(torch.cat([xs, xb], dim=1))
return out