amend

pytorch · Nov 22, 2023 · 82b668e · 82b668e
1 parent 9bf4f48
commit 82b668e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tensordict/base.py b/tensordict/base.py
@@ -3034,6 +3034,7 @@ def map(
             raise ValueError(f"Got incompatible dimension {dim_orig}")
 
         self_split = _split_tensordict(self, chunksize, num_chunks, num_workers, dim)
+        print('number of tds', len(self_split))
         chunksize = 1
         out = pool.imap(fn, self_split, chunksize)
         out = torch.cat(list(out), dim)