I’ve tried to rerun the code with “CUDA_LAUNCH_BLOCKING=” before python. Unfortunately, the code couldn’t run into the training process with this flag.
I’ve added  “torch.autograd.set_detect_anomaly(True)” into the code, then the error became:
 File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 539, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/linear.py", line 87, in forward
    return F.linear(input, self.weight, self.bias)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/functional.py", line 1365, in linear
    ret = torch.addmm(bias, input, weight.t())

Traceback (most recent call last):
  File "/workspace/pyface/engine/trainer.py", line 287, in <module>
    main(args)
  File "/workspace/pyface/engine/trainer.py", line 69, in main
    main_worker(args.gpu, ngpus_per_node, args)
  File "/workspace/pyface/engine/trainer.py", line 137, in main_worker
    do_train(train_loader, model, criterion, optimizer, epoch, args, architect, valid_loader)
  File "/workspace/pyface/engine/trainer.py", line 270, in do_train
    loss.backward()
  File "/opt/conda/lib/python3.6/site-packages/torch/tensor.py", line 118, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/opt/conda/lib/python3.6/site-packages/torch/autograd/__init__.py", line 93, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function 'AddmmBackward' returned nan values in its 2th output. 

It seems some nan in the grad?