From d05528913a51116459a000cb818b02b7c3c58de4 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 3 Jul 2024 15:16:16 +0100 Subject: [PATCH] amend --- benchmarks/distributed/dataloading.py | 3 +-- benchmarks/tensorclass/test_tensorclass_speed.py | 2 +- benchmarks/tensorclass/test_torch_functions.py | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/benchmarks/distributed/dataloading.py b/benchmarks/distributed/dataloading.py index 5d1777b57..0c8b03cdf 100644 --- a/benchmarks/distributed/dataloading.py +++ b/benchmarks/distributed/dataloading.py @@ -28,8 +28,7 @@ import torch import tqdm -from tensordict import MemoryMappedTensor -from tensordict.prototype import tensorclass +from tensordict import MemoryMappedTensor, tensorclass from torch import multiprocessing as mp, nn from torch.distributed import rpc from torch.utils.data import DataLoader diff --git a/benchmarks/tensorclass/test_tensorclass_speed.py b/benchmarks/tensorclass/test_tensorclass_speed.py index 0c8bf8c7d..452f6b86a 100644 --- a/benchmarks/tensorclass/test_tensorclass_speed.py +++ b/benchmarks/tensorclass/test_tensorclass_speed.py @@ -9,7 +9,7 @@ import pytest import torch -from tensordict.prototype import tensorclass +from tensordict import tensorclass @tensorclass diff --git a/benchmarks/tensorclass/test_torch_functions.py b/benchmarks/tensorclass/test_torch_functions.py index 0dc3c560c..7ed717872 100644 --- a/benchmarks/tensorclass/test_torch_functions.py +++ b/benchmarks/tensorclass/test_torch_functions.py @@ -6,7 +6,7 @@ import pytest import torch -from tensordict.prototype import tensorclass +from tensordict import tensorclass @tensorclass