From 31073a719880aa1126499bc06e7ef8f07f7a048c Mon Sep 17 00:00:00 2001
From: 15694147069 <591861959@qq.com>
Date: Thu, 7 Apr 2022 15:54:25 +0800
Subject: [PATCH] non_blocking=True

---
 .../3D_Nest_Unet_ID0476_for_PyTorch/train.py                  | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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 641e8cd07b..94468feb1e 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
-- 
Gitee