This is the full code:
import argparse
import logger
import time
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.optim as optim
import transforms
from data import MyDataset
from segnet import SegNet
from torch.autograd import Variable

def train(epoch):
    model.train()

    # update learning rate
    exp_lr_scheduler.step()

    total_loss = 0
    total_accuracy = 0

    # iteration over the batches
    for batch_idx, (img, gt) in enumerate(train_loader):

        if use_cuda:
            img = img.cuda(async=True)
            gt = gt.cuda(async=True)

        input = Variable(img)
        target = Variable(gt)

        # initialize gradients
        optimizer.zero_grad()

        # predictions
        output = model(input)

        """
        output is (24, 2, 224, 224)
        target is (24, 1, 224, 224)
        Here I change target.view() and type in order to use nn.CrossEntropyLoss()
        """
        
        tb = target.size(0)
        tc = target.size(1)
        th = target.size(2)
        tw = target.size(3)
        target_long = target.view(tb, th, tw).long()

        loss = cren_loss(output.cuda(), target_long.cuda())
        loss.backward()
        optimizer.step()

        """
        This is a segmentation task, so in the next part I compute how many 1 pixels are correctly classificated
        as 1 and how many 0 pixels are correctly 0. Then I simply calculate the mean of foreground and background
        accuracy.
        """
        
        output_pred = softmax(output)
        _, prediction = output_pred.max(dim=1)
        prediction = prediction.unsqueeze(1)

        mat_zero2zero = ((prediction == 0) * (target == 0)).int()
        mat_one2one = ((prediction == 1) * (target == 1)).int()

        prediction_back = mat_zero2zero.sum().float()
        target_back = target.numel() - target.sum()

        prediction_fore = mat_one2one.sum().float()
        target_fore = target.sum()

        acc_back = prediction_back / target_back
        acc_fore = prediction_fore / target_fore
        accuracy = (acc_back + acc_fore) / 2

        # TensorBoard logging
        info = {'train-loss': loss.data[0],
                'train-accuracy': accuracy}

        for tag, value in info.items():
            log.scalar_summary(tag, value, batch_idx + 1)

        print('batch: %5s | loss: %.3f | acc_back: %.3f | acc_fore: %.3f | acc: %.3f |'
              % (str(batch_idx + 1) + '/' + str(len(train_loader)),
                 loss.data[0],
                 acc_back,
                 acc_fore,
                 accuracy),
              time.strftime("%H:%M:%S", time.gmtime(time.time())),
              'training')

        total_loss += loss.data[0]
        total_accuracy += accuracy

    return total_loss / len(train_loader), total_accuracy / len(train_loader)


# training settings
parser = argparse.ArgumentParser(description='PyTorch SegNet')
parser.add_argument('--epochs', type=int, default=10, help='train epochs') 
parser.add_argument('--lr', type=float, default=0.0001, help='learning rate')
parser.add_argument('--momentum', type=float, default=0.5, help='SGD momentum')
parser.add_argument('--resume', '-r', action='store_true', help='resume from checkpoint')
args = parser.parse_args()

# cuda
use_cuda = torch.cuda.is_available()

input_nbr = 3
label_nbr = 2
img_size = 224

batch_size = 24
num_workers = 4

start_epoch = 0

softmax = torch.nn.Softmax(dim=1)

if use_cuda:
    cren_loss = nn.CrossEntropyLoss().cuda()
else:
    cren_loss = nn.CrossEntropyLoss()

# create SegNet model
model = SegNet(input_nbr, label_nbr)
model.load_from_filename('/path/to/pretrained/weights')

# convert to cuda if needed
if use_cuda:
    model.cuda()
    cudnn.benchmark = True
else:
    model.float()

# finetuning
ftparams = ['conv11d.weight', 'conv11d.bias']
for name, param in model.named_parameters():
    if name not in ftparams:
        param.requires_grad = False

# define the optimizer
optimizer = optim.SGD(model.conv11d.parameters(), lr=args.lr, momentum=args.momentum)
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)

# define data
root_dir_img = '/path/to/img/dir'
root_dir_gt = './path/to/gt/dir'

transform_train = transforms.Compose([
    transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),
    transforms.RandomResizedCrop(img_size),
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    transforms.ToTensor()
])

train_dataset = MyDataset(root_dir_img, root_dir_gt, transform_train)
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=num_workers,
    pin_memory=True
)

# Set the logger
log = logger.Logger('./logs')

for epoch in range(start_epoch, start_epoch + args.epochs):
    print('epoch: %5s' % str(epoch+1))

    # training
    train_loss, train_acc = train(epoch)
    print('\nepoch: %5s | loss: %.3f | acc: %.3f |'
          % (str(epoch + 1) + '/' + str(start_epoch + args.epochs),
             train_loss,
             train_acc),
          time.strftime("%H:%M:%S", time.gmtime(time.time())),
          'training')

    print('\n')