Skip to content

Commit ce25d72

Browse files
committed
Fix dltensor operator tests (#3984)
* Fix dltensor operator tests (missing imports) * Fix backend_impl shadowing Pipeline import Signed-off-by: Kamil Tokarski <[email protected]>
1 parent ac98939 commit ce25d72

File tree

2 files changed

+6
-4
lines changed

2 files changed

+6
-4
lines changed

dali/test/python/test_dltensor_operator.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
1+
# Copyright (c) 2019, 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.
@@ -27,13 +27,15 @@
2727
def setup_pytorch():
2828
global torch_dlpack
2929
global torch
30-
import torch as torch
30+
import torch
31+
import torch.utils.dlpack as torch_dlpack
3132
global torch_stream
3233
torch_stream = torch.cuda.Stream()
3334

3435

3536
def setup_mxnet():
3637
global mxnd
38+
from mxnet import ndarray as mxnd
3739

3840

3941
def setup_cupy():

dali/test/python/test_external_source_pytorch_dlpack.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
1+
# Copyright (c) 2020, 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.
@@ -16,8 +16,8 @@
1616
import nvidia.dali.fn as fn
1717
import nvidia.dali.ops as ops
1818
import torch
19-
from nvidia.dali import Pipeline
2019
from nvidia.dali.backend_impl import * # noqa: F401, F403
20+
from nvidia.dali import Pipeline
2121
from torch.utils.dlpack import to_dlpack, from_dlpack
2222

2323
from test_utils import check_output

0 commit comments

Comments
 (0)