mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix device error

Browse Source
This commit is contained in:
James Betker 2021-06-05 14:21:32 -06:00
parent 5f0cc65f3b
commit af52751d6b
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
codes/models/classifiers/cifar_resnet_branched.py
View File

@ -140,7 +140,7 @@ class ResNet(nn.Module):
for t in self.tails:
tailouts.append(t(output))
tailouts = torch.stack(tailouts, dim=0)
return (tailouts[coarse_label] * torch.eye(n=bs).view(bs,bs,1)).sum(dim=1)
return (tailouts[coarse_label] * torch.eye(n=bs, device=x.device).view(bs,bs,1)).sum(dim=1)
@register_model
def register_cifar_resnet18_branched(opt_net, opt):