Try to run the following script example in 0.4.0 and 0.3.0. In 0.3.0, this script runs normally and the returned loss_main is in torch.Size([1]). However, in 0.4.0, the returned loss_main is in torch.Size([]). Maybe 0.4.0 has some usage updates that I am not aware of.
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import torch.nn.functional as F
#from pycrayon import CrayonClient

def cross_entropy_3d(x, y, w=None):
  n, c, _, _, _ = x.size()
  # x_t : (N*X*Y*Z, C)
  x_t = torch.transpose(torch.transpose(torch.transpose(x, 1, 2), 2, 3), 3, 4).contiguous().view(-1, c)
  y_t = torch.transpose(torch.transpose(torch.transpose(y, 1, 2), 2, 3), 3, 4).contiguous().view(-1).long()
  loss = F.cross_entropy(x_t, y_t, weight=w)
  #loss = loss / x_t.size()[0]
  return loss


# output = torch.rand(5, 2, 3, 3, 3)
# label_batch = torch.randint(0, 2, (5, 1, 3, 3, 3), dtype=torch.long)
# output, label_batch = Variable(output.cuda()), Variable(label_batch.cuda())
output = np.random.randn(5, 2, 3, 3, 3)
label_batch = np.random.randint(2, size=(5, 1, 3, 3, 3))
output, label_batch = Variable(torch.from_numpy(output).cuda()), Variable(torch.from_numpy(label_batch).cuda())

loss_main = cross_entropy_3d(output, label_batch)
print(loss_main.shape)
print(loss_main.cpu().data.numpy()[0])