From 7cbc695227b4fded3003fa94fdc6277bd2a5fea9 Mon Sep 17 00:00:00 2001 From: kuizhiqing Date: Wed, 7 Dec 2022 11:07:33 +0800 Subject: [PATCH] use np for cpu tensor --- ppcls/arch/backbone/legendary_models/swin_transformer.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/ppcls/arch/backbone/legendary_models/swin_transformer.py b/ppcls/arch/backbone/legendary_models/swin_transformer.py index 728fb99d..3d83d0e8 100644 --- a/ppcls/arch/backbone/legendary_models/swin_transformer.py +++ b/ppcls/arch/backbone/legendary_models/swin_transformer.py @@ -64,10 +64,9 @@ class RollWithIndexSelect(paddle.autograd.PyLayer): def get_roll_index(H, W, shifts, place): - # following tensors will be created on cpu place with npu custom device - index = paddle.arange(0, H * W, dtype='int64').reshape([H, W]) # cpu - index_fp = paddle.roll(index, shifts=shifts, axis=(0, 1)).reshape([-1]) # cpu - index_bp = {i:idx for idx, i in enumerate(index_fp.numpy().tolist())} + index = np.arange(0, H * W, dtype=np.int64).reshape([H, W]) + index_fp = np.roll(index, shift=shifts, axis=(0, 1)).reshape([-1]) + index_bp = {i:idx for idx, i in enumerate(index_fp.tolist())} index_bp = [index_bp[i] for i in range(H * W)] index_fp = paddle.to_tensor(index_fp, place=place) index_bp = paddle.to_tensor(index_fp, dtype='int64', place=place) -- GitLab