From 5e44189d7eaf5e1ae25279fc3cf6402694eb8da2 Mon Sep 17 00:00:00 2001
From: D-X-Y <280835372@qq.com>
Date: Tue, 5 Nov 2019 23:38:53 +1100
Subject: [PATCH] fix bugs

---
 lib/datasets/SearchDatasetWrap.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/lib/datasets/SearchDatasetWrap.py b/lib/datasets/SearchDatasetWrap.py
index cb45a62..f7925b3 100644
--- a/lib/datasets/SearchDatasetWrap.py
+++ b/lib/datasets/SearchDatasetWrap.py
@@ -4,15 +4,18 @@ import torch.utils.data as data
 
 class SearchDataset(data.Dataset):
 
-  def __init__(self, name, data, train_split, valid_split):
+  def __init__(self, name, data, train_split, valid_split, check=True):
     self.datasetname = name
     self.data        = data
     self.train_split = train_split.copy()
     self.valid_split = valid_split.copy()
+    if check:
+      intersection = set(train_split).intersection(set(valid_split))
+      assert len(intersection) == 0, 'the splitted train and validation sets should have no intersection'
     self.length      = len(self.train_split)
 
   def __repr__(self):
-    return ('{name}(name={datasetname}, train={tr_L}, valid={val_L})'.format(name=self.__class__.__name__, tr_L=len(self.train_split), val_L=len(self.valid_split)))
+    return ('{name}(name={datasetname}, train={tr_L}, valid={val_L})'.format(name=self.__class__.__name__, datasetname=self.datasetname, tr_L=len(self.train_split), val_L=len(self.valid_split)))
 
   def __len__(self):
     return self.length