diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py
index 9679bbb1e6..601bd6f755 100644
--- a/deepmd/pt/infer/deep_eval.py
+++ b/deepmd/pt/infer/deep_eval.py
@@ -343,13 +343,17 @@ def _eval_model(
             natoms = len(atom_types[0])
 
         coord_input = torch.tensor(
-            coords.reshape([-1, natoms, 3]), dtype=GLOBAL_PT_FLOAT_PRECISION
-        ).to(DEVICE)
-        type_input = torch.tensor(atom_types, dtype=torch.long).to(DEVICE)
+            coords.reshape([-1, natoms, 3]),
+            dtype=GLOBAL_PT_FLOAT_PRECISION,
+            device=DEVICE,
+        )
+        type_input = torch.tensor(atom_types, dtype=torch.long, device=DEVICE)
         if cells is not None:
             box_input = torch.tensor(
-                cells.reshape([-1, 3, 3]), dtype=GLOBAL_PT_FLOAT_PRECISION
-            ).to(DEVICE)
+                cells.reshape([-1, 3, 3]),
+                dtype=GLOBAL_PT_FLOAT_PRECISION,
+                device=DEVICE,
+            )
         else:
             box_input = None
 
@@ -420,7 +424,7 @@ def eval_model(
         if cells is not None:
             assert isinstance(cells, torch.Tensor), err_msg
         assert isinstance(atom_types, torch.Tensor) or isinstance(atom_types, list)
-        atom_types = torch.tensor(atom_types, dtype=torch.long).to(DEVICE)
+        atom_types = torch.tensor(atom_types, dtype=torch.long, device=DEVICE)
     elif isinstance(coords, np.ndarray):
         if cells is not None:
             assert isinstance(cells, np.ndarray), err_msg
@@ -441,17 +445,17 @@ def eval_model(
         natoms = len(atom_types[0])
 
     coord_input = torch.tensor(
-        coords.reshape([-1, natoms, 3]), dtype=GLOBAL_PT_FLOAT_PRECISION
-    ).to(DEVICE)
-    type_input = torch.tensor(atom_types, dtype=torch.long).to(DEVICE)
+        coords.reshape([-1, natoms, 3]), dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE
+    )
+    type_input = torch.tensor(atom_types, dtype=torch.long, device=DEVICE)
     box_input = None
     if cells is None:
         pbc = False
     else:
         pbc = True
         box_input = torch.tensor(
-            cells.reshape([-1, 3, 3]), dtype=GLOBAL_PT_FLOAT_PRECISION
-        ).to(DEVICE)
+            cells.reshape([-1, 3, 3]), dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE
+        )
     num_iter = int((nframes + infer_batch_size - 1) / infer_batch_size)
 
     for ii in range(num_iter):
@@ -527,35 +531,37 @@ def eval_model(
         energy_out = (
             torch.cat(energy_out)
             if energy_out
-            else torch.zeros([nframes, 1], dtype=GLOBAL_PT_FLOAT_PRECISION).to(DEVICE)
+            else torch.zeros(
+                [nframes, 1], dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE
+            )
         )
         atomic_energy_out = (
             torch.cat(atomic_energy_out)
             if atomic_energy_out
-            else torch.zeros([nframes, natoms, 1], dtype=GLOBAL_PT_FLOAT_PRECISION).to(
-                DEVICE
+            else torch.zeros(
+                [nframes, natoms, 1], dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE
             )
         )
         force_out = (
             torch.cat(force_out)
             if force_out
-            else torch.zeros([nframes, natoms, 3], dtype=GLOBAL_PT_FLOAT_PRECISION).to(
-                DEVICE
+            else torch.zeros(
+                [nframes, natoms, 3], dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE
             )
         )
         virial_out = (
             torch.cat(virial_out)
             if virial_out
-            else torch.zeros([nframes, 3, 3], dtype=GLOBAL_PT_FLOAT_PRECISION).to(
-                DEVICE
+            else torch.zeros(
+                [nframes, 3, 3], dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE
             )
         )
         atomic_virial_out = (
             torch.cat(atomic_virial_out)
             if atomic_virial_out
             else torch.zeros(
-                [nframes, natoms, 3, 3], dtype=GLOBAL_PT_FLOAT_PRECISION
-            ).to(DEVICE)
+                [nframes, natoms, 3, 3], dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE
+            )
         )
         updated_coord_out = torch.cat(updated_coord_out) if updated_coord_out else None
         logits_out = torch.cat(logits_out) if logits_out else None
diff --git a/deepmd/pt/infer/inference.py b/deepmd/pt/infer/inference.py
index cce9291e3b..e97623dd24 100644
--- a/deepmd/pt/infer/inference.py
+++ b/deepmd/pt/infer/inference.py
@@ -171,7 +171,8 @@ def __init__(
 
     @staticmethod
     def get_data(data):
-        batch_data = next(iter(data))
+        with torch.device("cpu"):
+            batch_data = next(iter(data))
         for key in batch_data.keys():
             if key == "sid" or key == "fid":
                 continue
@@ -235,7 +236,8 @@ def run(self):
                 ),  # setting to 0 diverges the behavior of its iterator; should be >=1
                 drop_last=False,
             )
-            data = iter(dataloader)
+            with torch.device("cpu"):
+                data = iter(dataloader)
 
             single_results = {}
             sum_natoms = 0
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index b688362b73..29ca9c8f96 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -16,6 +16,9 @@
     FittingOutputDef,
     OutputVariableDef,
 )
+from deepmd.pt.utils import (
+    env,
+)
 from deepmd.pt.utils.nlist import (
     build_multiple_neighbor_list,
     get_multiple_nlist_key,
@@ -91,9 +94,17 @@ def get_model_sels(self) -> List[List[int]]:
 
     def _sort_rcuts_sels(self) -> Tuple[List[float], List[int]]:
         # sort the pair of rcut and sels in ascending order, first based on sel, then on rcut.
-        rcuts = torch.tensor(self.get_model_rcuts(), dtype=torch.float64)
-        nsels = torch.tensor(self.get_model_nsels())
-        zipped = torch.stack([torch.tensor(rcuts), torch.tensor(nsels)], dim=0).T
+        rcuts = torch.tensor(
+            self.get_model_rcuts(), dtype=torch.float64, device=env.DEVICE
+        )
+        nsels = torch.tensor(self.get_model_nsels(), device=env.DEVICE)
+        zipped = torch.stack(
+            [
+                torch.tensor(rcuts, device=env.DEVICE),
+                torch.tensor(nsels, device=env.DEVICE),
+            ],
+            dim=0,
+        ).T
         inner_sorting = torch.argsort(zipped[:, 1], dim=0)
         inner_sorted = zipped[inner_sorting]
         outer_sorting = torch.argsort(inner_sorted[:, 0], stable=True)
@@ -285,7 +296,7 @@ def __init__(
         self.smin_alpha = smin_alpha
 
         # this is a placeholder being updated in _compute_weight, to handle Jit attribute init error.
-        self.zbl_weight = torch.empty(0, dtype=torch.float64)
+        self.zbl_weight = torch.empty(0, dtype=torch.float64, device=env.DEVICE)
 
     def serialize(self) -> dict:
         return {
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index 4019fda423..c79e742d63 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -12,6 +12,9 @@
     FittingOutputDef,
     OutputVariableDef,
 )
+from deepmd.pt.utils import (
+    env,
+)
 from deepmd.utils.pair_tab import (
     PairTab,
 )
@@ -156,7 +159,7 @@ def forward_atomic(
         pairwise_rr = self._get_pairwise_dist(
             extended_coord, masked_nlist
         )  # (nframes, nloc, nnei)
-        self.tab_data = self.tab_data.view(
+        self.tab_data = self.tab_data.to(device=env.DEVICE).view(
             int(self.tab_info[-1]), int(self.tab_info[-1]), int(self.tab_info[2]), 4
         )
 
@@ -164,7 +167,8 @@ def forward_atomic(
         # i_type : (nframes, nloc), this is atype.
         # j_type : (nframes, nloc, nnei)
         j_type = extended_atype[
-            torch.arange(extended_atype.size(0))[:, None, None], masked_nlist
+            torch.arange(extended_atype.size(0), device=env.DEVICE)[:, None, None],
+            masked_nlist,
         ]
 
         raw_atomic_energy = self._pair_tabulated_inter(
diff --git a/deepmd/pt/model/descriptor/repformer_layer.py b/deepmd/pt/model/descriptor/repformer_layer.py
index 21ae0ff6f3..66ce38c0f7 100644
--- a/deepmd/pt/model/descriptor/repformer_layer.py
+++ b/deepmd/pt/model/descriptor/repformer_layer.py
@@ -326,7 +326,7 @@ def __init__(
         sel = [sel] if isinstance(sel, int) else sel
         self.nnei = sum(sel)
         assert len(sel) == 1
-        self.sel = torch.tensor(sel)
+        self.sel = torch.tensor(sel, device=env.DEVICE)
         self.sec = self.sel
         self.axis_dim = axis_dim
         self.set_davg_zero = set_davg_zero
diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py
index 0f8add1d8d..30682a4605 100644
--- a/deepmd/pt/model/descriptor/se_a.py
+++ b/deepmd/pt/model/descriptor/se_a.py
@@ -448,6 +448,7 @@ def forward(
             )  # shape is [nframes*nall, self.ndescrpt]
             xyz_scatter = torch.empty(
                 1,
+                device=env.DEVICE,
             )
             ret = self.filter_layers_old[0](dmatrix)
             xyz_scatter = ret
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index 19bc514a2d..9f0fd81842 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -240,8 +240,10 @@ def _format_nlist(
                         nlist,
                         -1
                         * torch.ones(
-                            [n_nf, n_nloc, nnei - n_nnei], dtype=nlist.dtype
-                        ).to(nlist.device),
+                            [n_nf, n_nloc, nnei - n_nnei],
+                            dtype=nlist.dtype,
+                            device=nlist.device,
+                        ),
                     ],
                     dim=-1,
                 )
diff --git a/deepmd/pt/model/network/network.py b/deepmd/pt/model/network/network.py
index 8b5b3cf998..9ef7b3366a 100644
--- a/deepmd/pt/model/network/network.py
+++ b/deepmd/pt/model/network/network.py
@@ -32,7 +32,7 @@
 
 
 def Tensor(*shape):
-    return torch.empty(shape, dtype=env.GLOBAL_PT_FLOAT_PRECISION)
+    return torch.empty(shape, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)
 
 
 class Dropout(nn.Module):
@@ -332,7 +332,13 @@ def __init__(
         bias: bool = True,
         init: str = "default",
     ):
-        super().__init__(d_in, d_out, bias=bias, dtype=env.GLOBAL_PT_FLOAT_PRECISION)
+        super().__init__(
+            d_in,
+            d_out,
+            bias=bias,
+            dtype=env.GLOBAL_PT_FLOAT_PRECISION,
+            device=env.DEVICE,
+        )
 
         self.use_bias = bias
 
@@ -552,6 +558,7 @@ def __init__(self, type_nums, embed_dim, bavg=0.0, stddev=1.0):
             embed_dim,
             padding_idx=type_nums,
             dtype=env.GLOBAL_PT_FLOAT_PRECISION,
+            device=env.DEVICE,
         )
         # nn.init.normal_(self.embedding.weight[:-1], mean=bavg, std=stddev)
 
@@ -799,7 +806,7 @@ def __init__(
             temperature=temperature,
         )
         self.attn_layer_norm = nn.LayerNorm(
-            self.embed_dim, dtype=env.GLOBAL_PT_FLOAT_PRECISION
+            self.embed_dim, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
         )
         if self.ffn:
             self.ffn_embed_dim = ffn_embed_dim
diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py
index d57d460f6d..ed2dfbc02b 100644
--- a/deepmd/pt/model/task/ener.py
+++ b/deepmd/pt/model/task/ener.py
@@ -268,7 +268,7 @@ def __init__(
             bias_atom_e = np.zeros([self.ntypes])
         if not use_tebd:
             assert self.ntypes == len(bias_atom_e), "Element count mismatches!"
-        bias_atom_e = torch.tensor(bias_atom_e)
+        bias_atom_e = torch.tensor(bias_atom_e, device=env.DEVICE)
         self.register_buffer("bias_atom_e", bias_atom_e)
 
         filter_layers_dipole = []
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index 8537be6e12..5a783e412b 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -156,7 +156,8 @@ def get_data_loader(_training_data, _validation_data, _training_params):
                 drop_last=False,
                 pin_memory=True,
             )
-            training_data_buffered = BufferedIterator(iter(training_dataloader))
+            with torch.device("cpu"):
+                training_data_buffered = BufferedIterator(iter(training_dataloader))
             validation_dataloader = DataLoader(
                 _validation_data,
                 sampler=valid_sampler,
@@ -166,7 +167,8 @@ def get_data_loader(_training_data, _validation_data, _training_params):
                 pin_memory=True,
             )
 
-            validation_data_buffered = BufferedIterator(iter(validation_dataloader))
+            with torch.device("cpu"):
+                validation_data_buffered = BufferedIterator(iter(validation_dataloader))
             if _training_params.get("validation_data", None) is not None:
                 valid_numb_batch = _training_params["validation_data"].get(
                     "numb_btch", 1
@@ -519,7 +521,8 @@ def step(_step_id, task_key="Default"):
                     if not torch.isfinite(grad_norm).all():
                         # check local gradnorm single GPU case, trigger NanDetector
                         raise FloatingPointError("gradients are Nan/Inf")
-                self.optimizer.step()
+                with torch.device("cpu"):
+                    self.optimizer.step()
                 self.scheduler.step()
             elif self.opt_type == "LKF":
                 if isinstance(self.loss, EnergyStdLoss):
diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py
index 7844d6d741..2125f9cdee 100644
--- a/deepmd/pt/utils/dataloader.py
+++ b/deepmd/pt/utils/dataloader.py
@@ -120,8 +120,9 @@ def construct_dataset(system):
             self.total_batch += len(system_dataloader)
         # Initialize iterator instances for DataLoader
         self.iters = []
-        for item in self.dataloaders:
-            self.iters.append(iter(item))
+        with torch.device("cpu"):
+            for item in self.dataloaders:
+                self.iters.append(iter(item))
 
     def set_noise(self, noise_settings):
         # noise_settings['noise_type'] # "trunc_normal", "normal", "uniform"
@@ -250,5 +251,6 @@ def get_weighted_sampler(training_data, prob_style, sys_prob=False):
     log.info("Generated weighted sampler with prob array: " + str(probs))
     # training_data.total_batch is the size of one epoch, you can increase it to avoid too many  rebuilding of iteraters
     len_sampler = training_data.total_batch * max(env.NUM_WORKERS, 1)
-    sampler = WeightedRandomSampler(probs, len_sampler, replacement=True)
+    with torch.device("cpu"):
+        sampler = WeightedRandomSampler(probs, len_sampler, replacement=True)
     return sampler
diff --git a/deepmd/pt/utils/nlist.py b/deepmd/pt/utils/nlist.py
index 43036604e2..0e2d9785f8 100644
--- a/deepmd/pt/utils/nlist.py
+++ b/deepmd/pt/utils/nlist.py
@@ -116,14 +116,14 @@ def build_neighbor_list(
         nlist = nlist[:, :, :nsel]
     else:
         rr = torch.cat(
-            [rr, torch.ones([batch_size, nloc, nsel - nnei]).to(rr.device) + rcut],
+            [rr, torch.ones([batch_size, nloc, nsel - nnei], device=rr.device) + rcut],
             dim=-1,
         )
         nlist = torch.cat(
             [
                 nlist,
-                torch.ones([batch_size, nloc, nsel - nnei], dtype=nlist.dtype).to(
-                    rr.device
+                torch.ones(
+                    [batch_size, nloc, nsel - nnei], dtype=nlist.dtype, device=rr.device
                 ),
             ],
             dim=-1,
@@ -289,7 +289,7 @@ def extend_coord_with_ghosts(
 
     """
     nf, nloc = atype.shape
-    aidx = torch.tile(torch.arange(nloc).unsqueeze(0), [nf, 1])
+    aidx = torch.tile(torch.arange(nloc, device=env.DEVICE).unsqueeze(0), [nf, 1])
     if cell is None:
         nall = nloc
         extend_coord = coord.clone()
diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py
index 051fddd14b..38f71d6994 100644
--- a/deepmd/pt/utils/stat.py
+++ b/deepmd/pt/utils/stat.py
@@ -32,16 +32,17 @@ def make_stat_input(datasets, dataloaders, nbatches):
     log.info(f"Packing data for statistics from {len(datasets)} systems")
     for i in range(len(datasets)):
         sys_stat = {key: [] for key in keys}
-        iterator = iter(dataloaders[i])
-        for _ in range(nbatches):
-            try:
-                stat_data = next(iterator)
-            except StopIteration:
-                iterator = iter(dataloaders[i])
-                stat_data = next(iterator)
-            for dd in stat_data:
-                if dd in keys:
-                    sys_stat[dd].append(stat_data[dd])
+        with torch.device("cpu"):
+            iterator = iter(dataloaders[i])
+            for _ in range(nbatches):
+                try:
+                    stat_data = next(iterator)
+                except StopIteration:
+                    iterator = iter(dataloaders[i])
+                    stat_data = next(iterator)
+                for dd in stat_data:
+                    if dd in keys:
+                        sys_stat[dd].append(stat_data[dd])
         for key in keys:
             if not isinstance(sys_stat[key][0], list):
                 if sys_stat[key][0] is None:
diff --git a/source/tests/pt/__init__.py b/source/tests/pt/__init__.py
index fdbdd73f79..1a6de0591a 100644
--- a/source/tests/pt/__init__.py
+++ b/source/tests/pt/__init__.py
@@ -3,3 +3,5 @@
 
 torch.set_num_threads(1)
 torch.set_num_interop_threads(1)
+# testing purposes; device should always be set explicitly
+torch.set_default_device("cuda:9999999")
diff --git a/source/tests/pt/model/test_autodiff.py b/source/tests/pt/model/test_autodiff.py
index e69e894af6..b1745b384e 100644
--- a/source/tests/pt/model/test_autodiff.py
+++ b/source/tests/pt/model/test_autodiff.py
@@ -54,9 +54,9 @@ def test(
         places = 8
         delta = 1e-5
         natoms = 5
-        cell = torch.rand([3, 3], dtype=dtype)
-        cell = (cell + cell.T) + 5.0 * torch.eye(3)
-        coord = torch.rand([natoms, 3], dtype=dtype)
+        cell = torch.rand([3, 3], dtype=dtype, device="cpu")
+        cell = (cell + cell.T) + 5.0 * torch.eye(3, device="cpu")
+        coord = torch.rand([natoms, 3], dtype=dtype, device="cpu")
         coord = torch.matmul(coord, cell)
         atype = torch.IntTensor([0, 0, 0, 1, 1])
         # assumes input to be numpy tensor
@@ -66,7 +66,10 @@ def np_infer(
             coord,
         ):
             e0, f0, v0 = eval_model(
-                self.model, torch.tensor(coord).unsqueeze(0), cell.unsqueeze(0), atype
+                self.model,
+                torch.tensor(coord, device=env.DEVICE).unsqueeze(0),
+                cell.unsqueeze(0),
+                atype,
             )
             ret = {
                 "energy": e0.squeeze(0),
@@ -92,9 +95,9 @@ def test(
         places = 8
         delta = 1e-4
         natoms = 5
-        cell = torch.rand([3, 3], dtype=dtype)
-        cell = (cell) + 5.0 * torch.eye(3)
-        coord = torch.rand([natoms, 3], dtype=dtype)
+        cell = torch.rand([3, 3], dtype=dtype, device="cpu")
+        cell = (cell) + 5.0 * torch.eye(3, device="cpu")
+        coord = torch.rand([natoms, 3], dtype=dtype, device="cpu")
         coord = torch.matmul(coord, cell)
         atype = torch.IntTensor([0, 0, 0, 1, 1])
         # assumes input to be numpy tensor
@@ -106,8 +109,10 @@ def np_infer(
         ):
             e0, f0, v0 = eval_model(
                 self.model,
-                torch.tensor(stretch_box(coord, cell, new_cell)).unsqueeze(0),
-                torch.tensor(new_cell).unsqueeze(0),
+                torch.tensor(
+                    stretch_box(coord, cell, new_cell), device="cpu"
+                ).unsqueeze(0),
+                torch.tensor(new_cell, device="cpu").unsqueeze(0),
                 atype,
             )
             ret = {
diff --git a/source/tests/pt/model/test_deeppot.py b/source/tests/pt/model/test_deeppot.py
index 3ef901d0e3..da2e554704 100644
--- a/source/tests/pt/model/test_deeppot.py
+++ b/source/tests/pt/model/test_deeppot.py
@@ -12,6 +12,7 @@
 )
 
 import numpy as np
+import torch
 
 from deepmd.infer.deep_pot import DeepPot as DeepPotUni
 from deepmd.pt.entrypoints.main import (
@@ -43,7 +44,8 @@ def setUp(self):
         trainer = get_trainer(deepcopy(self.config))
         trainer.run()
 
-        input_dict, label_dict, _ = trainer.get_data(is_train=False)
+        with torch.device("cpu"):
+            input_dict, label_dict, _ = trainer.get_data(is_train=False)
         trainer.wrapper(**input_dict, label=label_dict, cur_lr=1.0)
         self.model = "model.pt"
 
diff --git a/source/tests/pt/model/test_descriptor_dpa1.py b/source/tests/pt/model/test_descriptor_dpa1.py
index 07cb684afa..ced3bbac57 100644
--- a/source/tests/pt/model/test_descriptor_dpa1.py
+++ b/source/tests/pt/model/test_descriptor_dpa1.py
@@ -38,11 +38,9 @@ def setUp(self):
             8.178091365465004481e-01,
             6.159552512682983760e00,
         ]
-        self.cell = (
-            torch.tensor(cell, dtype=env.GLOBAL_PT_FLOAT_PRECISION)
-            .view(1, 3, 3)
-            .to(env.DEVICE)
-        )
+        self.cell = torch.tensor(
+            cell, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
+        ).view(1, 3, 3)
         coord = [
             2.978060152121375648e00,
             3.588469695887098077e00,
@@ -60,12 +58,12 @@ def setUp(self):
             6.575011420004332585e00,
             6.825240650611076099e00,
         ]
-        self.coord = (
-            torch.tensor(coord, dtype=env.GLOBAL_PT_FLOAT_PRECISION)
-            .view(1, -1, 3)
-            .to(env.DEVICE)
-        )
-        self.atype = torch.IntTensor([0, 0, 0, 1, 1]).view(1, -1).to(env.DEVICE)
+        self.coord = torch.tensor(
+            coord, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
+        ).view(1, -1, 3)
+        self.atype = torch.tensor(
+            [0, 0, 0, 1, 1], dtype=torch.int32, device=env.DEVICE
+        ).view(1, -1)
         self.ref_d = torch.tensor(
             [
                 8.382518544113587780e-03,
@@ -230,7 +228,8 @@ def setUp(self):
                 1.518237240362583541e-03,
             ],
             dtype=env.GLOBAL_PT_FLOAT_PRECISION,
-        ).to(env.DEVICE)
+            device=env.DEVICE,
+        )
         with open(Path(CUR_DIR) / "models" / "dpa1.json") as fp:
             self.model_json = json.load(fp)
         self.file_model_param = Path(CUR_DIR) / "models" / "dpa1.pth"
diff --git a/source/tests/pt/model/test_descriptor_dpa2.py b/source/tests/pt/model/test_descriptor_dpa2.py
index d94686fee1..96f734a7d8 100644
--- a/source/tests/pt/model/test_descriptor_dpa2.py
+++ b/source/tests/pt/model/test_descriptor_dpa2.py
@@ -40,11 +40,9 @@ def setUp(self):
             8.178091365465004481e-01,
             6.159552512682983760e00,
         ]
-        self.cell = (
-            torch.tensor(cell, dtype=env.GLOBAL_PT_FLOAT_PRECISION)
-            .view(1, 3, 3)
-            .to(env.DEVICE)
-        )
+        self.cell = torch.tensor(
+            cell, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
+        ).view(1, 3, 3)
         coord = [
             2.978060152121375648e00,
             3.588469695887098077e00,
@@ -62,12 +60,12 @@ def setUp(self):
             6.575011420004332585e00,
             6.825240650611076099e00,
         ]
-        self.coord = (
-            torch.tensor(coord, dtype=env.GLOBAL_PT_FLOAT_PRECISION)
-            .view(1, -1, 3)
-            .to(env.DEVICE)
-        )
-        self.atype = torch.IntTensor([0, 0, 0, 1, 1]).view(1, -1).to(env.DEVICE)
+        self.coord = torch.tensor(
+            coord, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
+        ).view(1, -1, 3)
+        self.atype = torch.tensor(
+            [0, 0, 0, 1, 1], dtype=torch.int32, device=env.DEVICE
+        ).view(1, -1)
         self.ref_d = torch.tensor(
             [
                 8.435412613327306630e-01,
@@ -112,7 +110,8 @@ def setUp(self):
                 -3.073744832541754762e-02,
             ],
             dtype=env.GLOBAL_PT_FLOAT_PRECISION,
-        ).to(env.DEVICE)
+            device=env.DEVICE,
+        )
         with open(Path(CUR_DIR) / "models" / "dpa2_hyb.json") as fp:
             self.model_json = json.load(fp)
         self.file_model_param = Path(CUR_DIR) / "models" / "dpa2.pth"
diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py
index 3f67043767..5ecfb72481 100644
--- a/source/tests/pt/model/test_dipole_fitting.py
+++ b/source/tests/pt/model/test_dipole_fitting.py
@@ -128,16 +128,16 @@ def setUp(self) -> None:
         self.rcut_smth = 0.5
         self.sel = [46, 92, 4]
         self.nf = 1
-        self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype).to(env.DEVICE)
-        self.shift = torch.tensor([4, 4, 4], dtype=dtype).to(env.DEVICE)
-        self.atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE)
+        self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype, device=env.DEVICE)
+        self.shift = torch.tensor([4, 4, 4], dtype=dtype, device=env.DEVICE)
+        self.atype = torch.tensor([0, 0, 0, 1, 1], dtype=torch.int32, device=env.DEVICE)
         self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE)
-        self.cell = torch.rand([3, 3], dtype=dtype).to(env.DEVICE)
-        self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3).to(env.DEVICE)
+        self.cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE)
+        self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3, device=env.DEVICE)
 
     def test_rot(self):
         atype = self.atype.reshape(1, 5)
-        rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype).to(env.DEVICE)
+        rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype, device=env.DEVICE)
         coord_rot = torch.matmul(self.coord, rmat)
         rng = np.random.default_rng()
         for mixed_types, nfp, nap in itertools.product(
diff --git a/source/tests/pt/model/test_embedding_net.py b/source/tests/pt/model/test_embedding_net.py
index c7782e61f2..87e8a97444 100644
--- a/source/tests/pt/model/test_embedding_net.py
+++ b/source/tests/pt/model/test_embedding_net.py
@@ -55,7 +55,7 @@ def get_single_batch(dataset, index=None):
     for key in ["coord", "box", "force", "energy", "virial", "atype", "natoms"]:
         if key in np_batch.keys():
             np_batch[key] = np.expand_dims(np_batch[key], axis=0)
-            pt_batch[key] = torch.as_tensor(np_batch[key])
+            pt_batch[key] = torch.as_tensor(np_batch[key], device=env.DEVICE)
     np_batch["coord"] = np_batch["coord"].reshape(1, -1)
     np_batch["natoms"] = np_batch["natoms"][0]
     np_batch["force"] = np_batch["force"].reshape(1, -1)
diff --git a/source/tests/pt/model/test_env_mat.py b/source/tests/pt/model/test_env_mat.py
index 0746f9e8df..ee262e7ee5 100644
--- a/source/tests/pt/model/test_env_mat.py
+++ b/source/tests/pt/model/test_env_mat.py
@@ -100,14 +100,14 @@ def test_consistency(
         em0 = EnvMat(self.rcut, self.rcut_smth)
         mm0, ww0 = em0.call(self.coord_ext, self.atype_ext, self.nlist, davg, dstd)
         mm1, _, ww1 = prod_env_mat_se_a(
-            torch.tensor(self.coord_ext, dtype=dtype),
-            torch.tensor(self.nlist, dtype=int),
-            torch.tensor(self.atype_ext[:, :nloc], dtype=int),
-            davg,
-            dstd,
+            torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE),
+            torch.tensor(self.nlist, dtype=int, device=env.DEVICE),
+            torch.tensor(self.atype_ext[:, :nloc], dtype=int, device=env.DEVICE),
+            torch.tensor(davg, device=env.DEVICE),
+            torch.tensor(dstd, device=env.DEVICE),
             self.rcut,
             self.rcut_smth,
         )
-        np.testing.assert_allclose(mm0, mm1)
-        np.testing.assert_allclose(ww0, ww1)
+        np.testing.assert_allclose(mm0, mm1.detach().cpu().numpy())
+        np.testing.assert_allclose(ww0, ww1.detach().cpu().numpy())
         np.testing.assert_allclose(mm0[0][self.perm[: self.nloc]], mm0[1])
diff --git a/source/tests/pt/model/test_linear_atomic_model.py b/source/tests/pt/model/test_linear_atomic_model.py
index dcfff0b387..aae1a66618 100644
--- a/source/tests/pt/model/test_linear_atomic_model.py
+++ b/source/tests/pt/model/test_linear_atomic_model.py
@@ -52,8 +52,8 @@ def test_pairwise(self, mock_loadtxt):
                 [0.25, 0.0, 0.0, 0.0],
             ]
         )
-        extended_atype = torch.tensor([[0, 0]]).to(env.DEVICE)
-        nlist = torch.tensor([[[1], [-1]]]).to(env.DEVICE)
+        extended_atype = torch.tensor([[0, 0]], device=env.DEVICE)
+        nlist = torch.tensor([[[1], [-1]]], device=env.DEVICE)
 
         ds = DescrptSeA(
             rcut=0.3,
@@ -85,8 +85,9 @@ def test_pairwise(self, mock_loadtxt):
                         [0.0, 0.0, 0.0],
                         [0.0, dist, 0.0],
                     ],
-                ]
-            ).to(env.DEVICE)
+                ],
+                device=env.DEVICE,
+            )
 
             wgt_model.forward_atomic(extended_coord, extended_atype, nlist)
 
@@ -106,7 +107,8 @@ def test_pairwise(self, mock_loadtxt):
                 [0.0, 0.0],
             ],
             dtype=torch.float64,
-        ).to(env.DEVICE)
+            device=env.DEVICE,
+        )
         torch.testing.assert_close(results, excepted_res, rtol=0.0001, atol=0.0001)
 
 
diff --git a/source/tests/pt/model/test_model.py b/source/tests/pt/model/test_model.py
index 789ad80ada..d8c7de39c3 100644
--- a/source/tests/pt/model/test_model.py
+++ b/source/tests/pt/model/test_model.py
@@ -337,7 +337,7 @@ def test_consistency(self):
         batch = my_ds.systems[0]._data_system._get_subdata(batch, 0)
         batch = my_ds.systems[0]._data_system.reformat_data_torch(batch)
         for key in ["coord", "atype", "box", "energy", "force"]:
-            batch[key] = torch.as_tensor(batch[key])
+            batch[key] = torch.as_tensor(batch[key], device=env.DEVICE)
             batch[key] = batch[key].unsqueeze(0)
         batch["coord"].requires_grad_(True)
         batch["natoms_vec"] = tmp
@@ -419,7 +419,7 @@ def step(step_id):
             var_name = torch2tf(name, last_layer_id=len(self.n_neuron))
             var_grad = vs_dict[var_name].gradient
             param_grad = param.grad.cpu()
-            var_grad = torch.tensor(var_grad)
+            var_grad = torch.tensor(var_grad, device="cpu")
             assert np.allclose(var_grad, param_grad, rtol=rtol, atol=atol)
 
 
diff --git a/source/tests/pt/model/test_nlist.py b/source/tests/pt/model/test_nlist.py
index 27c03acfaa..616af93081 100644
--- a/source/tests/pt/model/test_nlist.py
+++ b/source/tests/pt/model/test_nlist.py
@@ -26,12 +26,12 @@ def setUp(self):
         self.ns = 5 * 5 * 3
         self.nall = self.ns * self.nloc
         self.cell = torch.tensor(
-            [[1, 0, 0], [0.4, 0.8, 0], [0.1, 0.3, 2.1]], dtype=dtype
-        ).to(env.DEVICE)
-        self.icoord = torch.tensor([[0, 0, 0], [0.5, 0.5, 0.1]], dtype=dtype).to(
-            env.DEVICE
+            [[1, 0, 0], [0.4, 0.8, 0], [0.1, 0.3, 2.1]], dtype=dtype, device=env.DEVICE
         )
-        self.atype = torch.tensor([0, 1], dtype=torch.int).to(env.DEVICE)
+        self.icoord = torch.tensor(
+            [[0, 0, 0], [0.5, 0.5, 0.1]], dtype=dtype, device=env.DEVICE
+        )
+        self.atype = torch.tensor([0, 1], dtype=torch.int, device=env.DEVICE)
         [self.cell, self.icoord, self.atype] = [
             ii.unsqueeze(0) for ii in [self.cell, self.icoord, self.atype]
         ]
@@ -53,8 +53,9 @@ def setUp(self):
             [
                 [0, 0, 0, 0, 0, 0, -1, -1, -1, -1, 1, 1, 1, 1, -1, -1, -1, -1, -1, -1],
                 [0, 0, 0, 0, -1, -1, -1, -1, -1, -1, 1, 1, 1, 1, 1, 1, -1, -1, -1, -1],
-            ]
-        ).to(env.DEVICE)
+            ],
+            device=env.DEVICE,
+        )
 
     def test_build_notype(self):
         ecoord, eatype, mapping = extend_coord_with_ghosts(
@@ -174,39 +175,39 @@ def test_extend_coord(self):
         mm, cc = torch.unique(shift_vec[0][:, 0], dim=-1, return_counts=True)
         torch.testing.assert_close(
             mm,
-            torch.tensor([-2, -1, 0, 1, 2], dtype=dtype).to(env.DEVICE),
+            torch.tensor([-2, -1, 0, 1, 2], dtype=dtype, device=env.DEVICE),
             rtol=self.prec,
             atol=self.prec,
         )
         torch.testing.assert_close(
             cc,
-            torch.tensor([30, 30, 30, 30, 30], dtype=torch.long).to(env.DEVICE),
+            torch.tensor([30, 30, 30, 30, 30], dtype=torch.long, device=env.DEVICE),
             rtol=self.prec,
             atol=self.prec,
         )
         mm, cc = torch.unique(shift_vec[1][:, 1], dim=-1, return_counts=True)
         torch.testing.assert_close(
             mm,
-            torch.tensor([-2, -1, 0, 1, 2], dtype=dtype).to(env.DEVICE),
+            torch.tensor([-2, -1, 0, 1, 2], dtype=dtype, device=env.DEVICE),
             rtol=self.prec,
             atol=self.prec,
         )
         torch.testing.assert_close(
             cc,
-            torch.tensor([30, 30, 30, 30, 30], dtype=torch.long).to(env.DEVICE),
+            torch.tensor([30, 30, 30, 30, 30], dtype=torch.long, device=env.DEVICE),
             rtol=self.prec,
             atol=self.prec,
         )
         mm, cc = torch.unique(shift_vec[1][:, 2], dim=-1, return_counts=True)
         torch.testing.assert_close(
             mm,
-            torch.tensor([-1, 0, 1], dtype=dtype).to(env.DEVICE),
+            torch.tensor([-1, 0, 1], dtype=dtype, device=env.DEVICE),
             rtol=self.prec,
             atol=self.prec,
         )
         torch.testing.assert_close(
             cc,
-            torch.tensor([50, 50, 50], dtype=torch.long).to(env.DEVICE),
+            torch.tensor([50, 50, 50], dtype=torch.long, device=env.DEVICE),
             rtol=self.prec,
             atol=self.prec,
         )
diff --git a/source/tests/pt/model/test_pairtab_atomic_model.py b/source/tests/pt/model/test_pairtab_atomic_model.py
index 69e305b40b..021035e3de 100644
--- a/source/tests/pt/model/test_pairtab_atomic_model.py
+++ b/source/tests/pt/model/test_pairtab_atomic_model.py
@@ -11,6 +11,9 @@
 from deepmd.pt.model.atomic_model import (
     PairTabAtomicModel,
 )
+from deepmd.pt.utils import (
+    env,
+)
 from deepmd.pt.utils.utils import (
     to_numpy_array,
 )
@@ -45,21 +48,28 @@ def setUp(self, mock_loadtxt) -> None:
                     [0.01, 0.01, 0.02],
                     [0.05, 0.01, 0.01],
                 ],
-            ]
+            ],
+            device=env.DEVICE,
         )
 
         # nframes=2, nall=4
-        self.extended_atype = torch.tensor([[0, 1, 0, 1], [0, 0, 1, 1]])
+        self.extended_atype = torch.tensor(
+            [[0, 1, 0, 1], [0, 0, 1, 1]], device=env.DEVICE
+        )
 
         # nframes=2, nloc=2, nnei=2
-        self.nlist = torch.tensor([[[1, 2], [0, 2]], [[1, 2], [0, 3]]])
+        self.nlist = torch.tensor(
+            [[[1, 2], [0, 2]], [[1, 2], [0, 3]]], device=env.DEVICE
+        )
 
     def test_without_mask(self):
         result = self.model.forward_atomic(
             self.extended_coord, self.extended_atype, self.nlist
         )
         expected_result = torch.tensor(
-            [[[1.2000], [1.3614]], [[1.2000], [0.4000]]], dtype=torch.float64
+            [[[1.2000], [1.3614]], [[1.2000], [0.4000]]],
+            dtype=torch.float64,
+            device=env.DEVICE,
         )
 
         torch.testing.assert_close(
@@ -67,13 +77,17 @@ def test_without_mask(self):
         )
 
     def test_with_mask(self):
-        self.nlist = torch.tensor([[[1, -1], [0, 2]], [[1, 2], [0, 3]]])
+        self.nlist = torch.tensor(
+            [[[1, -1], [0, 2]], [[1, 2], [0, 3]]], device=env.DEVICE
+        )
 
         result = self.model.forward_atomic(
             self.extended_coord, self.extended_atype, self.nlist
         )
         expected_result = torch.tensor(
-            [[[0.8000], [1.3614]], [[1.2000], [0.4000]]], dtype=torch.float64
+            [[[0.8000], [1.3614]], [[1.2000], [0.4000]]],
+            dtype=torch.float64,
+            device=env.DEVICE,
         )
 
         torch.testing.assert_close(
@@ -91,7 +105,9 @@ def test_deserialize(self):
         torch.testing.assert_close(self.model.tab_data, model1.tab_data)
         torch.testing.assert_close(self.model.tab_info, model1.tab_info)
 
-        self.nlist = torch.tensor([[[1, -1], [0, 2]], [[1, 2], [0, 3]]])
+        self.nlist = torch.tensor(
+            [[[1, -1], [0, 2]], [[1, 2], [0, 3]]], device=env.DEVICE
+        )
         result = model1.forward_atomic(
             self.extended_coord, self.extended_atype, self.nlist
         )
@@ -116,12 +132,14 @@ def test_cross_deserialize(self):
 
         self.nlist = np.array([[[1, -1], [0, 2]], [[1, 2], [0, 3]]])
         result = model1.forward_atomic(
-            self.extended_coord.numpy(),
-            self.extended_atype.numpy(),
+            self.extended_coord.cpu().numpy(),
+            self.extended_atype.cpu().numpy(),
             self.nlist,
         )
         expected_result = self.model.forward_atomic(
-            self.extended_coord, self.extended_atype, torch.from_numpy(self.nlist)
+            self.extended_coord,
+            self.extended_atype,
+            torch.from_numpy(self.nlist).to(device=env.DEVICE),
         )
         np.testing.assert_allclose(
             result["energy"], to_numpy_array(expected_result["energy"]), 0.0001, 0.0001
@@ -159,10 +177,10 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None:
         )
 
         # nframes=1, nall=2
-        extended_atype = torch.tensor([[0, 0]])
+        extended_atype = torch.tensor([[0, 0]], device=env.DEVICE)
 
         # nframes=1, nloc=2, nnei=1
-        nlist = torch.tensor([[[1], [-1]]])
+        nlist = torch.tensor([[[1], [-1]]], device=env.DEVICE)
 
         results = []
 
@@ -206,7 +224,8 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None:
                         [0.0, 0.0, 0.0],
                         [0.0, dist, 0.0],
                     ],
-                ]
+                ],
+                device=env.DEVICE,
             )
 
             model = PairTabAtomicModel(tab_file=file_path, rcut=rcut, sel=2)
@@ -236,6 +255,7 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None:
                         ]
                     ],
                     dtype=torch.float64,
+                    device=env.DEVICE,
                 )
             ]
         ).reshape(14, 2)
diff --git a/source/tests/pt/model/test_permutation.py b/source/tests/pt/model/test_permutation.py
index 15359f873a..b97cb349ad 100644
--- a/source/tests/pt/model/test_permutation.py
+++ b/source/tests/pt/model/test_permutation.py
@@ -201,11 +201,11 @@ def test(
         self,
     ):
         natoms = 5
-        cell = torch.rand([3, 3], dtype=dtype).to(env.DEVICE)
-        cell = (cell + cell.T) + 5.0 * torch.eye(3).to(env.DEVICE)
-        coord = torch.rand([natoms, 3], dtype=dtype).to(env.DEVICE)
+        cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE)
+        cell = (cell + cell.T) + 5.0 * torch.eye(3, device=env.DEVICE)
+        coord = torch.rand([natoms, 3], dtype=dtype, device=env.DEVICE)
         coord = torch.matmul(coord, cell)
-        atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE)
+        atype = torch.tensor([0, 0, 0, 1, 1], dtype=torch.int32, device=env.DEVICE)
         idx_perm = [1, 0, 4, 3, 2]
         e0, f0, v0 = eval_model(
             self.model, coord.unsqueeze(0), cell.unsqueeze(0), atype
diff --git a/source/tests/pt/model/test_polarizability_fitting.py b/source/tests/pt/model/test_polarizability_fitting.py
index de43c57b8b..548219627b 100644
--- a/source/tests/pt/model/test_polarizability_fitting.py
+++ b/source/tests/pt/model/test_polarizability_fitting.py
@@ -149,17 +149,17 @@ def setUp(self) -> None:
         self.nf = 1
         self.nt = 3
         self.rng = np.random.default_rng()
-        self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype).to(env.DEVICE)
-        self.shift = torch.tensor([4, 4, 4], dtype=dtype).to(env.DEVICE)
-        self.atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE)
+        self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype, device=env.DEVICE)
+        self.shift = torch.tensor([4, 4, 4], dtype=dtype, device=env.DEVICE)
+        self.atype = torch.tensor([0, 0, 0, 1, 1], dtype=torch.int32, device=env.DEVICE)
         self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE)
-        self.cell = torch.rand([3, 3], dtype=dtype).to(env.DEVICE)
-        self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3).to(env.DEVICE)
+        self.cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE)
+        self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3, device=env.DEVICE)
         self.scale = self.rng.uniform(0, 1, self.nt).tolist()
 
     def test_rot(self):
         atype = self.atype.reshape(1, 5)
-        rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype).to(env.DEVICE)
+        rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype, device=env.DEVICE)
         coord_rot = torch.matmul(self.coord, rmat)
 
         for mixed_types, nfp, nap, fit_diag, scale in itertools.product(
diff --git a/source/tests/pt/model/test_region.py b/source/tests/pt/model/test_region.py
index e8a3346562..b06f4221fd 100644
--- a/source/tests/pt/model/test_region.py
+++ b/source/tests/pt/model/test_region.py
@@ -4,6 +4,9 @@
 import numpy as np
 import torch
 
+from deepmd.pt.utils import (
+    env,
+)
 from deepmd.pt.utils.preprocess import (
     Region3D,
 )
@@ -18,14 +21,14 @@
 class TestRegion(unittest.TestCase):
     def setUp(self):
         self.cell = torch.tensor(
-            [[1, 0, 0], [0.4, 0.8, 0], [0.1, 0.3, 2.1]], dtype=dtype
+            [[1, 0, 0], [0.4, 0.8, 0], [0.1, 0.3, 2.1]], dtype=dtype, device="cpu"
         )
         self.cell = self.cell.unsqueeze(0).unsqueeze(0)
         self.cell = torch.tile(self.cell, [4, 5, 1, 1])
         self.prec = 1e-8
 
     def test_inter_to_phys(self):
-        inter = torch.rand([4, 5, 3, 3], dtype=dtype)
+        inter = torch.rand([4, 5, 3, 3], dtype=dtype, device="cpu")
         phys = inter2phys(inter, self.cell)
         for ii in range(4):
             for jj in range(5):
@@ -45,7 +48,7 @@ def test_to_face_dist(self):
         dz = vol / sxy
         dy = vol / sxz
         dx = vol / syz
-        expected = torch.tensor([dx, dy, dz])
+        expected = torch.tensor([dx, dy, dz], device="cpu")
         dists = to_face_distance(self.cell)
         for ii in range(4):
             for jj in range(5):
@@ -57,19 +60,19 @@ def test_to_face_dist(self):
 class TestLegacyRegion(unittest.TestCase):
     def setUp(self):
         self.cell = torch.tensor(
-            [[1, 0, 0], [0.4, 0.8, 0], [0.1, 0.3, 2.1]], dtype=dtype
+            [[1, 0, 0], [0.4, 0.8, 0], [0.1, 0.3, 2.1]], dtype=dtype, device=env.DEVICE
         )
         self.prec = 1e-6
 
     def test_inter_to_phys(self):
-        inter = torch.rand([3, 3], dtype=dtype)
+        inter = torch.rand([3, 3], dtype=dtype, device=env.DEVICE)
         reg = Region3D(self.cell)
         phys = reg.inter2phys(inter)
         expected_phys = torch.matmul(inter, self.cell)
         torch.testing.assert_close(phys, expected_phys, rtol=self.prec, atol=self.prec)
 
     def test_inter_to_inter(self):
-        inter = torch.rand([3, 3], dtype=dtype)
+        inter = torch.rand([3, 3], dtype=dtype, device=env.DEVICE)
         reg = Region3D(self.cell)
         new_inter = reg.phys2inter(reg.inter2phys(inter))
         torch.testing.assert_close(inter, new_inter, rtol=self.prec, atol=self.prec)
diff --git a/source/tests/pt/model/test_rot.py b/source/tests/pt/model/test_rot.py
index 780d193ebd..0c3a34e2d5 100644
--- a/source/tests/pt/model/test_rot.py
+++ b/source/tests/pt/model/test_rot.py
@@ -32,15 +32,15 @@ def test(
     ):
         prec = 1e-10
         natoms = 5
-        cell = 10.0 * torch.eye(3, dtype=dtype).to(env.DEVICE)
-        coord = 2 * torch.rand([natoms, 3], dtype=dtype).to(env.DEVICE)
-        shift = torch.tensor([4, 4, 4], dtype=dtype).to(env.DEVICE)
-        atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE)
+        cell = 10.0 * torch.eye(3, dtype=dtype, device=env.DEVICE)
+        coord = 2 * torch.rand([natoms, 3], dtype=dtype, device=env.DEVICE)
+        shift = torch.tensor([4, 4, 4], dtype=dtype, device=env.DEVICE)
+        atype = torch.tensor([0, 0, 0, 1, 1], dtype=torch.int32, device=env.DEVICE)
         from scipy.stats import (
             special_ortho_group,
         )
 
-        rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype).to(env.DEVICE)
+        rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype, device=env.DEVICE)
 
         # rotate only coord and shift to the center of cell
         coord_rot = torch.matmul(coord, rmat)
@@ -74,11 +74,11 @@ def test(
 
         # rotate coord and cell
         torch.manual_seed(0)
-        cell = torch.rand([3, 3], dtype=dtype).to(env.DEVICE)
-        cell = (cell + cell.T) + 5.0 * torch.eye(3).to(env.DEVICE)
-        coord = torch.rand([natoms, 3], dtype=dtype).to(env.DEVICE)
+        cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE)
+        cell = (cell + cell.T) + 5.0 * torch.eye(3, device=env.DEVICE)
+        coord = torch.rand([natoms, 3], dtype=dtype, device=env.DEVICE)
         coord = torch.matmul(coord, cell)
-        atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE)
+        atype = torch.tensor([0, 0, 0, 1, 1], dtype=torch.int32, device=env.DEVICE)
         coord_rot = torch.matmul(coord, rmat)
         cell_rot = torch.matmul(cell, rmat)
         e0, f0, v0 = eval_model(
diff --git a/source/tests/pt/model/test_rotation.py b/source/tests/pt/model/test_rotation.py
index 18d7bb8553..caa6385c80 100644
--- a/source/tests/pt/model/test_rotation.py
+++ b/source/tests/pt/model/test_rotation.py
@@ -61,7 +61,7 @@ def get_rotation(self, index, rotation_matrix):
 def get_data(batch):
     inputs = {}
     for key in ["coord", "atype", "box"]:
-        inputs[key] = torch.as_tensor(batch[key])
+        inputs[key] = torch.as_tensor(batch[key], device=env.DEVICE)
         inputs[key] = inputs[key].unsqueeze(0).to(env.DEVICE)
     return inputs
 
@@ -74,7 +74,8 @@ def setUp(self):
         self.config["training"]["training_data"]["systems"] = data_file
         self.config["training"]["validation_data"]["systems"] = data_file
         self.rotation = special_ortho_group.rvs(3)
-        self.get_dataset(0)
+        with torch.device("cpu"):
+            self.get_dataset(0)
         self.get_model()
 
     def get_model(self):
diff --git a/source/tests/pt/model/test_saveload_dpa1.py b/source/tests/pt/model/test_saveload_dpa1.py
index 64229b8e9e..408afbef43 100644
--- a/source/tests/pt/model/test_saveload_dpa1.py
+++ b/source/tests/pt/model/test_saveload_dpa1.py
@@ -83,7 +83,8 @@ def setUp(self):
             drop_last=False,
             pin_memory=True,
         )
-        self.training_data = BufferedIterator(iter(self.training_dataloader))
+        with torch.device("cpu"):
+            self.training_data = BufferedIterator(iter(self.training_dataloader))
         self.loss = EnergyStdLoss(**self.config["loss"])
         self.cur_lr = 1
         self.task_key = "Default"
diff --git a/source/tests/pt/model/test_saveload_se_e2_a.py b/source/tests/pt/model/test_saveload_se_e2_a.py
index 0632e30b5b..382f119c30 100644
--- a/source/tests/pt/model/test_saveload_se_e2_a.py
+++ b/source/tests/pt/model/test_saveload_se_e2_a.py
@@ -83,7 +83,8 @@ def setUp(self):
             drop_last=False,
             pin_memory=True,
         )
-        self.training_data = BufferedIterator(iter(self.training_dataloader))
+        with torch.device("cpu"):
+            self.training_data = BufferedIterator(iter(self.training_dataloader))
         self.loss = EnergyStdLoss(**self.config["loss"])
         self.cur_lr = 1
         self.task_key = "Default"
diff --git a/source/tests/pt/model/test_smooth.py b/source/tests/pt/model/test_smooth.py
index fa9042a932..88d75a040c 100644
--- a/source/tests/pt/model/test_smooth.py
+++ b/source/tests/pt/model/test_smooth.py
@@ -37,27 +37,26 @@ def test(
         aprec = 1e-5 if self.aprec is None else self.aprec
 
         natoms = 10
-        cell = 8.6 * torch.eye(3, dtype=dtype).to(env.DEVICE)
-        atype = torch.randint(0, 3, [natoms])
-        coord0 = (
-            torch.tensor(
-                [
-                    0.0,
-                    0.0,
-                    0.0,
-                    4.0 - 0.5 * epsilon,
-                    0.0,
-                    0.0,
-                    0.0,
-                    4.0 - 0.5 * epsilon,
-                    0.0,
-                ],
-                dtype=dtype,
-            )
-            .view([-1, 3])
-            .to(env.DEVICE)
+        cell = 8.6 * torch.eye(3, dtype=dtype, device=env.DEVICE)
+        atype = torch.randint(0, 3, [natoms], device=env.DEVICE)
+        coord0 = torch.tensor(
+            [
+                0.0,
+                0.0,
+                0.0,
+                4.0 - 0.5 * epsilon,
+                0.0,
+                0.0,
+                0.0,
+                4.0 - 0.5 * epsilon,
+                0.0,
+            ],
+            dtype=dtype,
+            device=env.DEVICE,
+        ).view([-1, 3])
+        coord1 = torch.rand(
+            [natoms - coord0.shape[0], 3], dtype=dtype, device=env.DEVICE
         )
-        coord1 = torch.rand([natoms - coord0.shape[0], 3], dtype=dtype).to(env.DEVICE)
         coord1 = torch.matmul(coord1, cell)
         coord = torch.concat([coord0, coord1], dim=0)
 
diff --git a/source/tests/pt/model/test_trans.py b/source/tests/pt/model/test_trans.py
index a99d6c893f..23365f3c9a 100644
--- a/source/tests/pt/model/test_trans.py
+++ b/source/tests/pt/model/test_trans.py
@@ -31,12 +31,12 @@ def test(
         self,
     ):
         natoms = 5
-        cell = torch.rand([3, 3], dtype=dtype).to(env.DEVICE)
-        cell = (cell + cell.T) + 5.0 * torch.eye(3).to(env.DEVICE)
-        coord = torch.rand([natoms, 3], dtype=dtype).to(env.DEVICE)
+        cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE)
+        cell = (cell + cell.T) + 5.0 * torch.eye(3, device=env.DEVICE)
+        coord = torch.rand([natoms, 3], dtype=dtype, device=env.DEVICE)
         coord = torch.matmul(coord, cell)
-        atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE)
-        shift = (torch.rand([3], dtype=dtype) - 0.5).to(env.DEVICE) * 2.0
+        atype = torch.tensor([0, 0, 0, 1, 1], dtype=torch.int32, device=env.DEVICE)
+        shift = (torch.rand([3], dtype=dtype, device=env.DEVICE) - 0.5) * 2.0
         coord_s = torch.matmul(
             torch.remainder(torch.matmul(coord + shift, torch.linalg.inv(cell)), 1.0),
             cell,
diff --git a/source/tests/pt/model/test_unused_params.py b/source/tests/pt/model/test_unused_params.py
index f69d8ac835..c20a5f1dc5 100644
--- a/source/tests/pt/model/test_unused_params.py
+++ b/source/tests/pt/model/test_unused_params.py
@@ -58,9 +58,9 @@ def test_unused(self):
     def _test_unused(self, model_params):
         self.model = get_model(model_params).to(env.DEVICE)
         natoms = 5
-        cell = torch.rand([3, 3], dtype=dtype).to(env.DEVICE)
-        cell = (cell + cell.T) + 5.0 * torch.eye(3).to(env.DEVICE)
-        coord = torch.rand([natoms, 3], dtype=dtype).to(env.DEVICE)
+        cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE)
+        cell = (cell + cell.T) + 5.0 * torch.eye(3, device=env.DEVICE)
+        coord = torch.rand([natoms, 3], dtype=dtype, device=env.DEVICE)
         coord = torch.matmul(coord, cell)
         atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE)
         idx_perm = [1, 0, 4, 3, 2]
diff --git a/source/tests/pt/test_calculator.py b/source/tests/pt/test_calculator.py
index a35538250b..52b4b6cbbe 100644
--- a/source/tests/pt/test_calculator.py
+++ b/source/tests/pt/test_calculator.py
@@ -8,6 +8,7 @@
     Path,
 )
 
+import numpy as np
 import torch
 
 from deepmd.pt.entrypoints.main import (
@@ -39,7 +40,8 @@ def setUp(self):
         trainer = get_trainer(deepcopy(self.config))
         trainer.run()
 
-        input_dict, label_dict, _ = trainer.get_data(is_train=False)
+        with torch.device("cpu"):
+            input_dict, label_dict, _ = trainer.get_data(is_train=False)
         _, _, more_loss = trainer.wrapper(**input_dict, label=label_dict, cur_lr=1.0)
 
         self.calculator = DPCalculator("model.pt")
@@ -50,8 +52,8 @@ def test_calculator(self):
         )
 
         natoms = 5
-        cell = torch.eye(3, dtype=dtype) * 10
-        coord = torch.rand([natoms, 3], dtype=dtype)
+        cell = torch.eye(3, dtype=dtype, device="cpu") * 10
+        coord = torch.rand([natoms, 3], dtype=dtype, device="cpu")
         coord = torch.matmul(coord, cell)
         atype = torch.IntTensor([0, 0, 0, 1, 1])
         atomic_numbers = [1, 1, 1, 8, 8]
@@ -91,7 +93,7 @@ def test_calculator(self):
         assert isinstance(e0, float)
         assert f0.shape == (natoms, 3)
         assert v0.shape == (3, 3)
-        torch.testing.assert_close(e0, e1, rtol=low_prec, atol=prec)
-        torch.testing.assert_close(f0[idx_perm, :], f1, rtol=low_prec, atol=prec)
-        torch.testing.assert_close(s0, s1, rtol=low_prec, atol=prec)
-        torch.testing.assert_close(v0, v1, rtol=low_prec, atol=prec)
+        np.testing.assert_allclose(e0, e1, rtol=low_prec, atol=prec)
+        np.testing.assert_allclose(f0[idx_perm, :], f1, rtol=low_prec, atol=prec)
+        np.testing.assert_allclose(s0, s1, rtol=low_prec, atol=prec)
+        np.testing.assert_allclose(v0, v1, rtol=low_prec, atol=prec)
diff --git a/source/tests/pt/test_dp_test.py b/source/tests/pt/test_dp_test.py
index 3db66f073f..8d7dc9cd58 100644
--- a/source/tests/pt/test_dp_test.py
+++ b/source/tests/pt/test_dp_test.py
@@ -11,6 +11,7 @@
 )
 
 import numpy as np
+import torch
 
 from deepmd.pt.entrypoints.main import (
     get_trainer,
@@ -40,7 +41,8 @@ def test_dp_test(self):
         trainer = get_trainer(deepcopy(self.config))
         trainer.run()
 
-        input_dict, label_dict, _ = trainer.get_data(is_train=False)
+        with torch.device("cpu"):
+            input_dict, label_dict, _ = trainer.get_data(is_train=False)
         _, _, more_loss = trainer.wrapper(**input_dict, label=label_dict, cur_lr=1.0)
 
         tester = inference.Tester("model.pt", input_script=self.input_json)
diff --git a/source/tests/pt/test_sampler.py b/source/tests/pt/test_sampler.py
index 0ff16ed7c7..25980cc144 100644
--- a/source/tests/pt/test_sampler.py
+++ b/source/tests/pt/test_sampler.py
@@ -7,6 +7,7 @@
 )
 
 import numpy as np
+import torch
 from torch.utils.data import (
     DataLoader,
 )
@@ -69,7 +70,8 @@ def test_sampler_debug_info(self):
             drop_last=False,
             pin_memory=True,
         )
-        batch_data = next(iter(dataloader))
+        with torch.device("cpu"):
+            batch_data = next(iter(dataloader))
         sid = batch_data["sid"]
         fid = batch_data["fid"][0]
         coord = batch_data["coord"].squeeze(0)
diff --git a/source/tests/pt/test_stat.py b/source/tests/pt/test_stat.py
index 5cf6a953cc..1e3c707d6f 100644
--- a/source/tests/pt/test_stat.py
+++ b/source/tests/pt/test_stat.py
@@ -131,9 +131,9 @@ def my_merge(energy, natoms):
             natoms_lst = []
             for i in range(len(energy)):
                 for j in range(len(energy[i])):
-                    energy_lst.append(torch.tensor(energy[i][j]))
+                    energy_lst.append(torch.tensor(energy[i][j], device="cpu"))
                     natoms_lst.append(
-                        torch.tensor(natoms[i][j])
+                        torch.tensor(natoms[i][j], device="cpu")
                         .unsqueeze(0)
                         .expand(energy[i][j].shape[0], -1)
                     )