@ptrblck looks like @mariosfourn was correct while saying

torch.utils.data.random_split()  returns  idx  as  torch.Tensor rather than a float.

As per the example in question, indexing ants_dataset would work correctly but an error would be raised if accessing an index for train_dataset.
This could be resolved by adding idx = idx.item() but this would make indexing ants_dataset not functional.
A quick hack would be to have
    def __getitem__(self, idx):
        try:
          idx = idx.item()
        except:
          idx = idx
        return self.data.iloc[idx]

This way we’ll be able to access both ants_dataset as well as train_dataset using indices.