Fix: Correct concatenation of datasets in train conditioning

Summary: ChainDataset is iterable, and it toes not go along with a custom batch sampler.

Reviewed By: bottler

Differential Revision: D42742315

fbshipit-source-id: 40a715c8d24abe72cb2777634247d7467f628564

pytorch3d/implicitron/dataset/data_loader_map_provider.py

@@ -12,7 +12,7 @@
 from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
 from torch.utils.data import (
     BatchSampler,
-    ChainDataset,
+    ConcatDataset,
     DataLoader,
     RandomSampler,
     Sampler,
@@ -482,7 +482,7 @@ def _train_loader(
             num_batches=num_batches,
         )
         return DataLoader(
-            ChainDataset([dataset, train_dataset]),
+            ConcatDataset([dataset, train_dataset]),
             batch_sampler=sampler,
             **data_loader_kwargs,
         )