split_data_target issue

[jfulem]

Hi, I'm using a Dataset class (extending torch.utils.data.Dataset) which contains this method:

def __getitem__(self, idx) -> tuple[list[torch.Tensor], torch.Tensor]

Unfortunately, the split_data_target in torchensemble.utils.io can't handle that.

   if len(element) == 2:
        # Dataloader with one input and one target
        data, target = element[0], element[1]
        return [data.to(device)], target.to(device)  

maybe this modification would be useful:

    if len(element) == 2:
        # Dataloader with one input and one target
        data, target = element[0], element[1]
        if isinstance(data, list) or isinstance(data, tuple):
            return [d.to(device) for d in data], target.to(device)
        return [data.to(device)], target.to(device)# tensor -> list

[xuyxu]

def getitem(self, idx) -> tuple[list[torch.Tensor], torch.Tensor]

Hi @jfulem, is this specification of dataset generally used? Will appreciate your suggestions very much.

[jfulem]

Yes, getting this tuple[list[torch.Tensor], torch.Tensor] as the dataset item is very common.

[xuyxu]

Marked as a new feature, are you willing to work on this @jfulem ?

[jfulem]

@xuyxu, sure I can make the PR.