Traning procedure stucked at the end of one epoch

This error will happen on pytorch version greater than 0.4.1. I followed this instruction to build 0.4.1 from source. Weirdly, there is no problem with the pip installed 0.4.1 version.

My pytorch version:

In [1]: import torch

In [2]: torch.__version__
Out[2]: '0.5.0a0+a24163a'

My code is heavily borrowed from pytorch-CycleGAN-and-pix2pix. I use torch.nn.DataParallel to warp the module.

    # create dataset
    # create dataloader
    # create model
    for epoch in range(1, opt.niter):
        for i, data in enumerate(dataset):
            ...
            model.set_input(data)
            model.optimize_parameters()
            ...
        model.save(...)

I have tried different dataloader implemention.

  1. pytorch-CycleGAN-and-pix2pix implemention

  2. vanilla version

def create_dataloader(dataset_opt):
    dataset_class = find_dataset_using_name(dataset_opt.dataset_mode)
    dataset = dataset_class(dataset_opt)
    print("dataset [%s] was created" % type(dataset).__name__)
    return torch.utils.data.DataLoader(
        dataset,
        batch_size=dataset_opt.batchSize,
        shuffle=dataset_opt.shuffle_data,
        num_workers=dataset_opt.nThreads,
        sampler=None,
        drop_last=True,
        pin_memory=False), len(dataset)