From 0b1af4ea19a6f85fb1c9bf6bb8589490d70497bc Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Mon, 15 Aug 2022 18:33:19 +0800 Subject: [PATCH] Update test Signed-off-by: Jael Gu --- test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test.py b/test.py index 62ab8a4..1a3419d 100644 --- a/test.py +++ b/test.py @@ -7,6 +7,7 @@ import numpy import onnx import onnxruntime +device = 'cuda' if torch.cuda.is_available() else 'cpu' def to_numpy(tensor): return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy() @@ -14,7 +15,7 @@ def to_numpy(tensor): # decode = ops.audio_decode.ffmpeg() # audio = [x[0] for x in decode('path/to/audio.wav')] -audio = torch.rand(10, 256, 32) +audio = torch.rand(10, 256, 32).to(device) op = ops.audio_embedding.nnfp() out0 = op.get_op().model(audio)