[BUG] PT: display_if_exist is blocking

[njzjz]

Bug summary

The profiler shows that cudaStreamSynchronize happens in display_if_exist.

In display_if_exist, find_property is expected to be float. However, it is tensor(1., device='cuda:0'), a float32 tensor on the GPU, causing the synchronization.

deepmd-kit/deepmd/pt/loss/loss.py

Lines 32 to 43 in 0c0878e

	@staticmethod
	def display_if_exist(loss: torch.Tensor, find_property: float) -> torch.Tensor:
	"""Display NaN if labeled property is not found.
	Parameters
	----------
	loss : torch.Tensor
	the loss tensor
	find_property : float
	whether the property is found
	"""
	return loss if bool(find_property) else torch.nan

DeePMD-kit Version

0c0878e

Backend and its version

PyTorch 2.3.1

How did you download the software?

Built from source

Input Files, Running Commands, Error Log, etc.

Use examples/water/se_atten_compressible to debug.

Steps to Reproduce

cd examples/water/se_atten_compressible
dp --pt train input.json

Further Information, Files, and Links

No response