ValueError: probabilities contain NaN

I add apex in my code as follows:

But when I run it , problems occured as follows:
File "train.py", line 336, in
main()
File "train.py", line 167, in main
pbar=pbar)
File "train.py", line 214, in train
y_pred_raw, feature_matrix, attention_map = net(X)
File "/root/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 547, in call
result = self.forward(*input, **kwargs)
File "/root/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 152, in forward
outputs = self.parallel_apply(replicas, inputs, kwargs)
File "/root/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 162, in parallel_apply
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
File "/root/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 85, in parallel_apply
output.reraise()
File "/root/anaconda3/lib/python3.7/site-packages/torch/_utils.py", line 369, in reraise
raise self.exc_type(msg)
ValueError: Caught ValueError in replica 0 on device 0.
Original Traceback (most recent call last):
File "/root/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 60, in _worker
output = module(*input, **kwargs)
File "/root/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 547, in call
result = self.forward(*input, **kwargs)
File "/data/zhaoyanmei/projects/WS-DAN.PyTorch/models/wsdan.py", line 120, in forward
k_index = np.random.choice(self.M, 2, p=attention_weights.cpu().numpy())
File "mtrand.pyx", line 1144, in mtrand.RandomState.choice
ValueError: probabilities contain NaN

The traceback points to the external train.py and models/wsdan.py files, specifically the np.random.choice call receiving attention_weights containing NaN. Start by reproducing the failure and determining whether Apex is involved; done requires a minimal reproducible case or clear evidence that the problem lies outside Apex.