Thank you @ptrblck, this error is occasionally triggered during model inference, not often. If I encounter this exception during large-scale data inference task, how can I accurately find the wrong batch of data? As your said, CUDA operations are asynchronously, if I catch the exception and log the bad batch, can I locate this wrong batch?
like this:
try:
	batch = inference_queue.get(block=True)
	with torch.no_grad():
		input_ids = torch.tensor(batch['input_ids'],
								 dtype=torch.long).to(device)
		attention_mask = torch.tensor(batch['attention_mask'],
									  dtype=torch.long).to(device)
		inputs = {
			"input_ids": input_ids,
			"attention_mask": attention_mask
		}
		logits = model(**inputs)[0]
		probs = softmax(logits)
	probs = probs.cpu().numpy()
except Exception as e:
	logger.error(f'{traceback.format_exc()}, DoInference is dead!')
	with open('./BadBatch.pkl', 'wb') as f:
		pickle.dump(batch, f)
	logger.error(f'Bad batch recorded!')