albanD:

torch.autograd.set_detect_anomaly(True)


Thanks ! I add the  torch.autograd.set_detect_anomaly(True) and then get the error as follows:

/pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:57: UserWarning: Traceback of forward call that caused the error:
  File "/home/lth/anaconda3/lib/python3.7/threading.py", line 890, in _bootstrap
    self._bootstrap_inner()
  File "/home/lth/anaconda3/lib/python3.7/threading.py", line 926, in _bootstrap_inner
    self.run()
  File "/home/lth/anaconda3/lib/python3.7/threading.py", line 870, in run
    self._target(*self._args, **self._kwargs)
  File "/home/lth/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 60, in _worker
    output = module(*input, **kwargs)
  File "/home/lth/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 547, in __call__
    result = self.forward(*input, **kwargs)
  File "demo.py", line 58, in forward
    x = sum([tensor * w for tensor, w in zip(tensors, probs_i)])
  File "demo.py", line 58, in <listcomp>
    x = sum([tensor * w for tensor, w in zip(tensors, probs_i)])

Traceback (most recent call last):
  File "demo.py", line 210, in <module>
    loss.backward()
  File "/home/lth/anaconda3/lib/python3.7/site-packages/torch/tensor.py", line 118, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/lth/anaconda3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 93, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function MulBackward0 returned an invalid gradient at index 1 - expected device cuda:0 but got cuda:3

The error happens at  forward  of  EnergyConstrainedADMM_P_imagenet, but this error just happens when using torch.nn.DataParallel.
Looking forward to your reply, thanks !