fix ut

Author: caic99

source/tests/pt/model/test_saveload_dpa1.py

@@ -115,7 +115,8 @@ def create_wrapper(self, read: bool):
         return ModelWrapper(model, self.loss)
 
     def get_data(self):
-        batch_data = next(self.training_data)
+        with torch.device("cpu"):
+            batch_data = next(self.training_data)
         input_dict = {}
         for item in ["coord", "atype", "box"]:
             if item in batch_data:

source/tests/pt/model/test_saveload_se_e2_a.py

@@ -109,7 +109,8 @@ def create_wrapper(self):
         return ModelWrapper(model, self.loss)
 
     def get_data(self):
-        batch_data = next(self.training_data)
+        with torch.device("cpu"):
+            batch_data = next(self.training_data)
         input_dict = {}
         for item in ["coord", "atype", "box"]:
             if item in batch_data: