From ce0e3fb5dadc1d686884d073087250ea9028b1a8 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Tue, 5 Mar 2024 22:23:31 -0500
Subject: [PATCH] avoid the behavior changes

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 deepmd/pt/loss/tensor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py
index 2436b3bc6e..e0cf229e74 100644
--- a/deepmd/pt/loss/tensor.py
+++ b/deepmd/pt/loss/tensor.py
@@ -85,7 +85,7 @@ def forward(self, model_pred, label, natoms, learning_rate=0.0, mae=False):
         del learning_rate, mae
         loss = torch.tensor.zeros(
             1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
-        )
+        )[0]
         more_loss = {}
         if (
             self.has_local_weight