torch.Size报错：ValueError: Using a target size (torch.Size([128])) that is different to the input

问题

/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
1193 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194 return forward_call(*input, **kwargs)
1195 # Do not call functions when jit is used
1196 full_backward_hooks, non_full_backward_hooks = [], []

/lib/python3.7/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
617
…
-> 3088 “Please ensure they have the same size.”.format(target.size(), input.size())
3089 )
3090

ValueError: Using a target size (torch.Size([128])) that is different to the input size (torch.Size([256])) is deprecated. Please ensure they have the same size.
出现torch.Size不匹配的问题，并且这个不匹配还是整数倍，比如在我的例子中，target.size()是128，但input.size()是256.

分析

由于我是调用别人写好的库，不应该有这个问题，并且在github上也没有没有找到一样的问题，最后我看了下源码，这个库支持多GPU训练，我使用了2个GPU，于是我修改了源码改成单个GPU训练，问题解决。

解决办法

训练时制定单个GPU即可解决问题

os.environ["CUDA_VISIBLE_DEVICES"] = "0"