Is there any update for this?
I have the similar problem. I added flush=True to print statement and found that the hang actually happened in loss.backward(). But the strange thing is that it didn’t happen at the first iteration.
My training code:
    for epoch in range(args.epoch):
        log('EPOCH %d' % epoch)
        if args.distributed:
            train_sampler.set_epoch(epoch)
        for i, data in enumerate(trainloader):
            log('ITER %d' % i)
            net.train()
            inputs, labels = Variable(data[0]), Variable(data[1])
            if args.use_gpu:
                inputs, labels = inputs.cuda(), labels.cuda()
            log('ITER %d DATA LOADED' % i)

            outputs = net.forward(inputs)
            loss = criterion(outputs, labels)
            log('ITER %d FORWARDED' % i)

            optimizer.zero_grad()
            log('ITER %d ZERO_GRAD' % i)

            loss.backward()
            log('ITER %d BACKWARDED' % i)

            optimizer.step()
            log('ITER %d STEP' % i)

The output is:
2018-04-17_10:19:18 EPOCH 0
2018-04-17_10:19:21 ITER 0
2018-04-17_10:19:21 ITER 0 DATA LOADED
2018-04-17_10:19:30 ITER 0 FORWARDED
2018-04-17_10:19:30 ITER 0 ZERO_GRAD
2018-04-17_10:19:32 ITER 0 BACKWARDED
2018-04-17_10:19:32 ITER 0 STEP
2018-04-17_10:19:32 ITER 1
2018-04-17_10:19:32 ITER 1 DATA LOADED
2018-04-17_10:19:33 ITER 1 FORWARDED
2018-04-17_10:19:33 ITER 1 ZERO_GRAD