Hi Ptrblck, I’m using your loss as below, but it turns out

TypeError: softmax() received an invalid combination of arguments - got (Tensor), but expected one of: * (Tensor input, name dim, torch.dtype dtype)* (Tensor input, int dim, torch.dtype dtype)

my task is maybe the same as rajibmon, I changed the unet model to

Class UNet(nn.Module):

def __init__(self, in_channels=1,num_class =2,init_features=32):
    super(UNet, self).__init__()

    features = init_features
    self.num_classes = num_class
    # self.input = FCN._block(in_channels, features, name="input")
    # self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
    self.encoder1 = UNet._block(in_channels, features * 1, name="enc1")
    self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
    self.encoder2 = UNet._block(features * 1, features * 2, name="enc2")
    self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
    self.encoder3 = UNet._block(features * 2, features * 4, name="enc3")
    self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
    self.encoder4 = UNet._block(features * 4, features * 8, name='enc4')
    self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)

    self.bottleneck = UNet._block(features * 8, features * 16, name='bottleneck')


    self.upconv4 = nn.ConvTranspose2d(features * 16, features * 8, kernel_size=2, stride=2)
    self.decoder4 = UNet._block((features * 8)*2, features * 8, name="dec4")
    self.upconv3 = nn.ConvTranspose2d(features * 8, features * 4, kernel_size=2, stride=2)
    self.decoder3 = UNet._block((features * 4)*2, features * 4, name="dec3")
    self.upconv2 = nn.ConvTranspose2d(features * 4, features * 2, kernel_size=2, stride=2)
    self.decoder2 = UNet._block((features * 2)*2, features * 2, name="dec2")
    self.upconv1 = nn.ConvTranspose2d(features * 2, features, kernel_size=2, stride=2)
    self.decoder1 = UNet._block(features*2, features, name="dec1")

    self.conv = nn.Conv2d(in_channels=features, out_channels=num_class, kernel_size=1)
def forward(self, x):
    enc1 = self.encoder1(x)
    enc2 = self.encoder2(self.pool1(enc1))
    enc3 = self.encoder3(self.pool2(enc2))
    enc4 = self.encoder4(self.pool3(enc3))
    bottleneck  = self.bottleneck(self.pool4(enc4))
    dec4 = self.upconv4(bottleneck)
    dec4 = torch.cat((dec4,enc4),dim =1)
    dec4 = F.dropout(self.decoder4(dec4),0.5,training=True)
    dec3 = self.upconv3(dec4)
    dec3 = torch.cat((dec3,enc3),dim=1)
    dec3 = F.dropout(self.decoder3(dec3),0.5,training=True)
    dec2 = self.upconv2(dec3)
    dec2 = torch.cat((dec2,enc2),dim=1)
    dec2 = F.dropout(self.decoder2(dec2),0.5,training=True)
    dec1 = self.upconv1(dec2)
    dec1 = torch.cat((dec1,enc1),dim=1)
    dec1 = F.dropout(self.decoder1(dec1),0.5,training=True)
    out = self.conv(dec1)
    out = torch.softmax(out)
    return out

loss is now defined as

def weightedpixelcros(bseg,bgt,bs):
out = F.log_softmax(bseg)
weighs = torch.FloatTensor(bs,1,256,256).random_(1,3)
out = out.gather(1, bgt.view(bs, 1, 256,256))
wout = (outweighs).view(bs,-1)
wgloss = wout.sum(1) / weighs.view(bs, -1).sum(1)
wgloss = -1.0wgloss.mean()
return wgloss


net = UNet(in_channels=1, num_class=2, init_features=64)
net.to(device)
bone_seg = net(input_mr)
optimizer.zero_grad()
loss = weightedpixelcros(bone_gt, bone_seg, bs=BATCH_SIZE)
loss.backward()

Could you please help me to have a look if I’m doing something wrong??