Update unet_model

unet_model

@@ -71,4 +71,4 @@ class UNet(nn.Module):
         u1 = torch.cat([u1, c1], dim=1)
         d1 = self.dec1(u1)
 
-        return tor
+        return torch.sigmoid(self.conv_last(d1))