Support 0<data_fraction<1 for CustomDatasetWithoutLabels (#328)

* Support 0<data_fraction<1 for CustomDatasetWithoutLabels

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: kaland313

solo/data/pretrain_dataloader.py

@@ -352,16 +352,23 @@ def prepare_datasets(
 
     if data_fraction > 0:
         assert data_fraction < 1, "Only use data_fraction for values smaller than 1."
-        data = train_dataset.samples
-        files = [f for f, _ in data]
-        labels = [l for _, l in data]
-
         from sklearn.model_selection import train_test_split
 
-        files, _, labels, _ = train_test_split(
-            files, labels, train_size=data_fraction, stratify=labels, random_state=42
-        )
-        train_dataset.samples = [tuple(p) for p in zip(files, labels)]
+        if isinstance(train_dataset, CustomDatasetWithoutLabels):
+            files = train_dataset.images
+            (
+                files,
+                _,
+            ) = train_test_split(files, train_size=data_fraction, random_state=42)
+            train_dataset.images = files
+        else:
+            data = train_dataset.samples
+            files = [f for f, _ in data]
+            labels = [l for _, l in data]
+            files, _, labels, _ = train_test_split(
+                files, labels, train_size=data_fraction, stratify=labels, random_state=42
+            )
+            train_dataset.samples = [tuple(p) for p in zip(files, labels)]
 
     return train_dataset