there is an error in line 125

[Qiaohao501]

Traceback (most recent call last):
File "train.py", line 125, in
output_masked[:,d,:,:] = (input_mask[:,d,:,:].unsqueeze(1)) * output
RuntimeError: expand(torch.cuda.FloatTensor{[36, 1, 128, 128]}, size=[36, 128, 128]): the number of sizes provided (3) must be greater or equal to the number of dimensions in the tensor (4)

[Aut0matas]

origin: output_masked[:,d,:,:] = input_mask[:,d,:,:].unsqueeze(1) * output
fix: output_masked[:,d,:,:] = torch.squeeze(input_mask[:,d,:,:].unsqueeze(1) * output)