diff --git a/PyTorch/dev/cv/image_classification/3D_Nest_Unet_ID0476_for_PyTorch/train.py b/PyTorch/dev/cv/image_classification/3D_Nest_Unet_ID0476_for_PyTorch/train.py
index 641e8cd07b73f59a8e487d11457befaa59c8f300..94468feb1e9f56fa25912a96a9c8af377b844e04 100644
--- a/PyTorch/dev/cv/image_classification/3D_Nest_Unet_ID0476_for_PyTorch/train.py
+++ b/PyTorch/dev/cv/image_classification/3D_Nest_Unet_ID0476_for_PyTorch/train.py
@@ -158,8 +158,8 @@ def train(config, train_loader, model, criterion, optimizer,):
 
     pbar = tqdm(total=len(train_loader))
     for input, target, _ in train_loader:
-        input = input.npu()
-        target = target.npu()
+        input = input.npu(non_blocking=True)
+        target = target.npu(non_blocking=True)
         start_time=time.time()
         batchsize=input.shape[0]
         # compute output