DataLoader not iterable error

When I tried to use LightningDataModule, I first defined a DataModule like this:

class MultiSceneDataModuleNew(pl.LightningDataModule):
    def __init__(self, args, config):
        self.data_root = config.DATASET.DATA_ROOT  # Root path to EPIC-FIELDS dataset
        self.batch_size = args.batch_size
        self.num_workers = args.num_workers
        self.pin_memory = getattr(args, "pin_memory", True)

    def setup(self, stage: str = None):
        if stage == "fit" or stage is None:
            self.train_dataset = self._setup_dataset(mode="train")
            self.val_dataset = self._setup_dataset(mode="val")
  "Train & Val Dataset loaded!")

        elif stage == "test" or stage is None:
            self.test_dataset = self._setup_dataset(mode="test")
  "Test Dataset loaded!")

    def _setup_dataset(self, mode):
        datasets = []
        for scene_folder in tqdm(os.listdir(self.data_root), desc=f"Loading {mode} datasets"):
            scene_path = os.path.join(self.data_root, scene_folder)
            if os.path.isdir(scene_path):
                json_path = os.path.join(self.data_root, f"{scene_folder}.json")
                if os.path.exists(json_path):
                    datasets.append(EpicKitchensDataset(scene_path, json_path, mode))
        return ConcatDataset(datasets)

    def train_dataloader(self):
        dataloader = DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, pin_memory=self.pin_memory)
        return dataloader

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=False, pin_memory=self.pin_memory)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size, shuffle=False, pin_memory=self.pin_memory)

I also defined a Dataset Class which implements __len__() and __getitem__() methods, like this:

class EpicKitchensDataset(Dataset):
    def __init__(self, scene_folder, json_file, mode="train"):
        self.scene_folder = scene_folder
        with open(json_file, 'r') as f:
   = json.load(f)

        self.cam_matrices, self.world_to_cam_poses = self.parse_camera_data(['camera'])
        self.transform = transforms.ToTensor()
        self.depth_pipe = pipeline(task="depth-estimation", model="LiheYoung/depth-anything-small-hf")

    def parse_camera_data(self, camera_data):
        # Extract and construct intrinsic and extrinsic matrices
        fx, fy, cx, cy = camera_data['params'][:4]
        intrinsic_matrix = torch.tensor([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=torch.float32)

        world_to_cam_poses = []
        for pose in['images'].values():
            R = qvec2rotmat(pose[:4])
            t = np.array(pose[4:]).reshape(3, 1)
            extrinsic_matrix = np.hstack((R, t))
            extrinsic_matrix = np.vstack((extrinsic_matrix, [0, 0, 0, 1]))
            world_to_cam_poses.append(torch.tensor(extrinsic_matrix, dtype=torch.float32))

        return intrinsic_matrix, torch.stack(world_to_cam_poses)

    def __len__(self):
        return len(['images'])

    def __getitem__(self, idx):
        frame_name = list(['images'].keys())[idx]
        image_path = os.path.join(self.scene_folder, frame_name)
        # img = plt.imread(image_path)
        # img_tensor = self.transform(img) if img.dtype == np.float32 else self.transform(img / 255.0)
        # Use to load the image directly as a tensor
        if not os.path.exists(image_path):
            print(f"Image not found: {image_path}")
            img =
        # img =

        # Apply transformations: resize and convert to tensor
        img = self.transform(img)
        if img.shape[1:] != (256, 456):
            print(f"Resizing image {image_path} from {img.shape[1:]} to (256, 456)")
            img = resize(img, [256, 456])

        depth = self.depth_pipe(image_path)["depth"]
        depth_tensor = self.transform(depth)

        sample = {
            "img": img,
            "gt_depth": depth_tensor,
            "intrinsics": self.cam_matrices,
            "world_to_cam_pose": self.world_to_cam_poses[idx]
        return sample

The datamodule works fine when I called like this:

    data_module = MultiSceneDataModuleNew(args, config)
    print(next(iter(data_module.train_dataloader()))), datamodule=data_module)

However, I received the error below when I ran, datamodule=data_module):

Traceback (most recent call last):
  File "/python3.9/site-packages/lightning/pytorch/trainer/connectors/", line 402, in _check_dataloader_iterable
    iter(dataloader)  # type: ignore[call-overload]
TypeError: 'MultiSceneDataModuleNew' object is not iterable

During handling of the above exception, another exception occurred:
Traceback (most recent call last):
 File "/python3.9/site-packages/lightning/pytorch/trainer/connectors/", line 407, in _check_dataloader_iterable
    raise TypeError(
TypeError: An invalid dataloader was passed to ``. Found <src.lightning.data_new.MultiSceneDataModuleNew object at 0x7fe006450640>.
/python3.9/multiprocessing/ UserWarning: resource_tracker: process died unexpectedly, relaunching.  Some resources might leak.
  warnings.warn('resource_tracker: process died unexpectedly, '

I think the error was due to some error in the data module, but I cannot locate it. Could anyone give me some insight? Any help would be appreciated. Thanks! :slight_smile:

:face_with_head_bandage: :yum: Can anyone help?