Below code will create a mask for x
class get_mask(nn.Module):
  def __int__(self, in_c):
    self.conv1 = nn.Conv2d(in_c, 1 , 3, 1)
    self.norm = nn.BatchNorm2d(2)

  def forward(self, x):
    p = 1
    p2d =(p, p, p, p)
    x = self.conv1(F.pad(x,p2d))
    x = self.normalize(x)
    return F.sigmoid(x)

How modify the network to output a binary mask(zeros and ones) instead of F.sigmoid(x)?