1 files changed, 2 insertions, 3 deletions

diff --git a/data/csv.py b/data/csv.py
index 6525e45..d400757 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -3,7 +3,6 @@ import torch
 import json
 import numpy as np
 from pathlib import Path
-import pytorch_lightning as pl
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader, random_split
 from torchvision import transforms
@@ -42,7 +41,7 @@ class CSVDataItem(NamedTuple):
     nprompt: str
 
 
-class CSVDataModule(pl.LightningDataModule):
+class CSVDataModule():
     def __init__(
             self,
             batch_size: int,
@@ -141,7 +140,7 @@ class CSVDataModule(pl.LightningDataModule):
             items = [
                 item
                 for item in items
-                if "mode" in item and self.mode in item["mode"]
+                if "mode" in item and self.mode in item["mode"].split(", ")
             ]
         items = self.prepare_items(template, expansions, items)
         items = self.filter_items(items)