So this is my training code:
vgg_model = VGGNet(requires_grad=True)
fcn_model = FCNs(pretrained_net=vgg_model, n_class=2)
fcn_model = fcn_model.cuda()
criterion = nn.BCELoss().cuda()
optimizer = optim.SGD(fcn_model.parameters(), lr=1e-2, momentum=0.7)
saving_index = 0
i = 0
iou = 0
iou_bg = 0
for epo in range(30):
  saving_index +=1
  index = 0
  epo_loss = 0
  for item in train_loader:
      index += 1
      start = time.time()
      input = item['A']
      y = item['B']
      input = torch.autograd.Variable(input)
      y = torch.autograd.Variable(y)

      input = input.float()
      input = input.cuda()
      y = y.cuda()

      optimizer.zero_grad()
      output = fcn_model(input)
      output = nn.functional.sigmoid(output)
      loss = criterion(output, y)
      loss.backward()
      iter_loss = loss.item()
      epo_loss += iter_loss
      optimizer.step()

      output_np = output.cpu().data.numpy().copy()
      output_bg_np = np.argmax(output_np, axis=1)
      output_np = np.argmin(output_np, axis=1)
      
      y_np = y.cpu().data.numpy().copy()
      y_np = np.argmin(y_np, axis=1)
      input = input.cpu().data.numpy().copy().squeeze(1)
      y_not = np.where(y_np[0, :, :] == 1, 0, 1)


      intersect = (output_np[0, :, :]*y_np[0, :, :]).sum()
      union = (output_np[0, :, :] + y_np[0, :, :]).sum()
      intersect_bg = (output_bg_np*y_not).sum()
      union_bg = (output_bg_np + y_not).sum()
      iu = intersect / (union-intersect+0.00001)
      iu_bg = intersect_bg / (union_bg-intersect_bg+0.00001)

      fig, ax = plt.subplots(1,2, figsize=(10,3))  # 1 row, 2 columns
      ax[0].imshow(output_np[0, :, :])
      ax[1].imshow(y_np[0, :, :])
      iou = iu + iou
      iou_bg = iu_bg + iou_bg

      i = i + 1
      print("IOU is:", iou/i, "bg is: ", iou_bg/i)
      print("mean IOU is:", (iou/i+iou_bg/i)/2)

      if np.mod(index, 5) ==1:
          print('epoch {}, {}/{}, loss is {}'.format(epo, index, len(input), iter_loss))
          fig, ax = plt.subplots(1,2, figsize=(10,3))  # 1 row, 2 columns
          ax[0].imshow(output_np[0, :, :])
          ax[1].imshow(y_np[0, :, :])
  print('epoch loss = %f'%(epo_loss/len(train_loader)))
  
  if np.mod(saving_index, 5)==1:
      torch.save(fcn_model, './fcn_model_{}.pt'.format(epo))
      print('saveing checkpoints/fcn_model_{}.pt'.format(epo))
print("IOU is:", iou/i, "bg is: ", iou_bg/i)
print("mean IOU is:", (iou/i+iou_bg/i)/2)