amend

pytorch · Jul 3, 2024 · d055289 · d055289
1 parent 18b5c45
commit d055289
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 4 deletions.
diff --git 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
@@ -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
@@ -6,7 +6,7 @@
 import pytest
 import torch
 
-from tensordict.prototype import tensorclass
+from tensordict import tensorclass
 
 
 @tensorclass